diff --git a/.github/workflows/arrow-flight-tests.yml b/.github/workflows/arrow-flight-tests.yml index f6f39252bed78..22ebf9439fcb2 100644 --- a/.github/workflows/arrow-flight-tests.yml +++ b/.github/workflows/arrow-flight-tests.yml @@ -51,7 +51,7 @@ jobs: # Run Maven tests for the target module, excluding native tests - name: Maven Tests - run: ./mvnw test ${MAVEN_TEST} -pl ${{ matrix.modules }} -Dtest="*,!TestArrowFlightNativeQueries" + run: ./mvnw test ${MAVEN_TEST} -pl ${{ matrix.modules }} -Dtest="*,!TestArrowFlightNativeQueries*" prestocpp-linux-build-for-test: runs-on: ubuntu-22.04 @@ -207,7 +207,7 @@ jobs: mvn test \ ${MAVEN_TEST} \ -pl ${{ matrix.modules }} \ - -Dtest="TestArrowFlightNativeQueries" \ + -Dtest="TestArrowFlightNativeQueries*" \ -DPRESTO_SERVER=${PRESTO_SERVER_PATH} \ -DDATA_DIR=${RUNNER_TEMP} \ -Duser.timezone=America/Bahia_Banderas \ diff --git a/.github/workflows/presto-release-publish.yml b/.github/workflows/presto-release-publish.yml index 25d7d3af2f7cf..78e156eeea97a 100644 --- a/.github/workflows/presto-release-publish.yml +++ b/.github/workflows/presto-release-publish.yml @@ -19,6 +19,11 @@ on: type: boolean default: true required: false + publish_github_release: + description: 'Publish executable jars to github release' + type: boolean + default: true + required: false publish_image: description: 'Publish presto docker image' type: boolean @@ -33,11 +38,6 @@ on: description: 'prestissimo dependency image(e.g., prestodb/presto-native-dependency:latest)' required: false default: '' - publish_native_dep_image: - description: 'Publish prestissimo dependency docker image' - type: boolean - default: true - required: false tag_image_as_latest: description: 'Tag the published images as latest' type: boolean @@ -102,8 +102,12 @@ jobs: publish-maven-artifacts: needs: publish-release-tag if: | - (!failure() &&!cancelled()) && - (github.event.inputs.publish_maven == 'true' || github.event.inputs.publish_image == 'true' || github.event.inputs.publish_docs == 'true') + (!failure() &&!cancelled()) && ( + github.event.inputs.publish_maven == 'true' || + github.event.inputs.publish_image == 'true' || + github.event.inputs.publish_docs == 'true' || + github.event.inputs.publish_github_release == 'true' + ) runs-on: ubuntu-latest environment: release timeout-minutes: 60 @@ -181,18 +185,8 @@ jobs: EOL - - name: Maven build - run: mvn clean install -DskipTests - - - name: Upload artifacts - uses: actions/upload-artifact@v4 - with: - name: presto-artifacts-${{ env.RELEASE_TAG }} - retention-days: 1 - path: | - presto-server/target/presto-server-*.tar.gz - presto-cli/target/presto-cli-*-executable.jar - presto-docs/target/presto-docs-*.zip + - name: Build release artifacts + run: ./mvnw clean install -DskipTests - name: Set up GPG key env: @@ -213,28 +207,67 @@ jobs: GPG_PASSPHRASE: "${{ secrets.GPG_PASSPHRASE }}" run: | unset MAVEN_CONFIG - modules=$(find . -maxdepth 1 -type d \( ! -name . \) -exec test -f '{}/pom.xml' \; -print | sed 's|^./||') - for module in $modules; do - if [ "$module" = "presto-test-coverage" ]; then - echo "Skipping $module" - continue - fi - MODULE_PATH=$(awk -F'[<>]' '/^ / {print $3; exit}' "$module/pom.xml" | sed 's|\.|/|g' | sed 's|^com/facebook/presto||') - echo "Checking https://repo1.maven.org/maven2/com/facebook/presto${MODULE_PATH}/${module}/${RELEASE_TAG}" - if ! curl --silent --fail "https://repo1.maven.org/maven2/com/facebook/presto${MODULE_PATH}/${module}/${RELEASE_TAG}" > /dev/null; then - echo "Publishing ${module}" - ./mvnw -s ${{ github.workspace }}/settings.xml -V -B -U -e -T1C deploy \ - -Dgpg.passphrase="${{ secrets.GPG_PASSPHRASE }}" \ - -Dmaven.wagon.http.retryHandler.count=8 \ - -DskipTests \ - -Drelease.autoPublish=${{ env.MAVEN_AUTO_PUBLISH }} \ - -Poss-release \ - -Pdeploy-to-ossrh \ - -pl "$module" - else - echo "${module} ${RELEASE_TAG} already published, skipping." - fi - done + ./mvnw -s ${{ github.workspace }}/settings.xml -V -B -U -e -T1C deploy \ + -Dgpg.passphrase="${{ secrets.GPG_PASSPHRASE }}" \ + -Dmaven.wagon.http.retryHandler.count=8 \ + -DskipTests \ + -DstagingProfileId="${{ secrets.MAVEN_STAGING_PROFILE_ID }}" \ + -DkeepStagingRepositoryOnFailure=true \ + -DkeepStagingRepositoryOnCloseRuleFailure=true \ + -DautoReleaseAfterClose=true \ + -DstagingProgressTimeoutMinutes=60 \ + -DdeploymentName=presto-${{ env.RELEASE_TAG }} \ + -Drelease.autoPublish=${{ env.MAVEN_AUTO_PUBLISH }} \ + -DskipExecutableJar \ + -Poss-release \ + -Pdeploy-to-ossrh \ + -pl '!presto-test-coverage' + + - name: Upload artifacts + uses: actions/upload-artifact@v4 + with: + name: presto-artifacts-${{ env.RELEASE_TAG }} + retention-days: 7 + path: | + presto-server/target/presto-server-*.tar.gz + presto-docs/target/presto-docs-*.zip + presto-cli/target/presto-cli-*-executable.jar + presto-benchmark-driver/target/presto-benchmark-driver-*-executable.jar + presto-testing-server-launcher/target/presto-testing-server-launcher-*-executable.jar + + publish-github-release: + needs: publish-maven-artifacts + if: (!failure() && !cancelled()) && github.event.inputs.publish_github_release == 'true' + runs-on: ubuntu-latest + environment: release + timeout-minutes: 150 + permissions: + packages: write + contents: read + steps: + - name: Checkout + uses: actions/checkout@v4 + with: + ref: ${{ env.RELEASE_TAG}} + + - name: Download artifacts + uses: actions/download-artifact@v4 + with: + name: presto-artifacts-${{ env.RELEASE_TAG }} + path: ./ + + - name: Create github release + uses: softprops/action-gh-release@v2 + with: + tag_name: ${{ env.RELEASE_TAG }} + name: Presto ${{ env.RELEASE_TAG }} + token: ${{ secrets.PRESTODB_CI_TOKEN }} + body: | + See the release notes at https://prestodb.io/docs/current/release/release-${{ env.RELEASE_TAG }}.html + files: | + ./presto-cli/target/presto-cli-*-executable.jar + ./presto-benchmark-driver/target/presto-benchmark-driver-*-executable.jar + ./presto-testing-server-launcher/target/presto-testing-server-launcher-*-executable.jar publish-docker-image: needs: publish-maven-artifacts @@ -349,7 +382,7 @@ jobs: else echo "Building new depedency image" docker compose build centos-native-dependency - if [[ "${{ github.event.inputs.publish_native_dep_image }}" == "true" ]]; then + if [[ "${{ github.event.inputs.publish_native_image }}" == "true" ]]; then docker tag presto/prestissimo-dependency:centos9 ${{ github.repository_owner }}/presto-native-dependency:${{ env.RELEASE_TAG }}-${{ env.COMMIT_SHA }} docker push ${{ github.repository_owner }}/presto-native-dependency:${{ env.RELEASE_TAG }}-${{ env.COMMIT_SHA }} @@ -365,7 +398,17 @@ jobs: working-directory: presto-native-execution run: | df -h - docker compose build centos-native-runtime + docker compose build --build-arg EXTRA_CMAKE_FLAGS=" + -DPRESTO_ENABLE_PARQUET=ON \ + -DPRESTO_ENABLE_REMOTE_FUNCTIONS=ON \ + -DPRESTO_ENABLE_JWT=ON \ + -DPRESTO_STATS_REPORTER_TYPE=PROMETHEUS \ + -DPRESTO_MEMORY_CHECKER_TYPE=LINUX_MEMORY_CHECKER \ + -DPRESTO_ENABLE_SPATIAL=ON \ + -DPRESTO_ENABLE_TESTING=OFF \ + -DPRESTO_ENABLE_S3=ON" \ + --build-arg NUM_THREADS=2 \ + centos-native-runtime - name: Add release tag working-directory: presto-native-execution diff --git a/.github/workflows/prestocpp-linux-build-and-unit-test.yml b/.github/workflows/prestocpp-linux-build-and-unit-test.yml index b28f2292252ea..7917102a6b06b 100644 --- a/.github/workflows/prestocpp-linux-build-and-unit-test.yml +++ b/.github/workflows/prestocpp-linux-build-and-unit-test.yml @@ -312,6 +312,84 @@ jobs: -Duser.timezone=America/Bahia_Banderas \ -T1C + prestocpp-linux-presto-on-spark-e2e-tests: + needs: prestocpp-linux-build-for-test + runs-on: ubuntu-22.04 + strategy: + fail-fast: false + matrix: + storage-format: [ "PARQUET", "DWRF" ] + enable-sidecar: [ "true", "false" ] + container: + image: prestodb/presto-native-dependency:0.293-20250522140509-484b00e + env: + MAVEN_OPTS: "-Xmx4G -XX:+ExitOnOutOfMemoryError" + MAVEN_FAST_INSTALL: "-B -V --quiet -T 1C -DskipTests -Dair.check.skip-all -Dmaven.javadoc.skip=true" + MAVEN_TEST: "-B -Dair.check.skip-all -Dmaven.javadoc.skip=true -DLogTestDurationListener.enabled=true --fail-at-end" + steps: + - uses: actions/checkout@v4 + + - name: Fix git permissions + # Usually actions/checkout does this but as we run in a container + # it doesn't work + run: git config --global --add safe.directory ${GITHUB_WORKSPACE} + + - name: Download artifacts + uses: actions/download-artifact@v4 + with: + name: presto-native-build + path: presto-native-execution/_build/release + + # Permissions are lost when uploading. Details here: https://github.com/actions/upload-artifact/issues/38 + - name: Restore execute permissions and library path + run: | + chmod +x ${GITHUB_WORKSPACE}/presto-native-execution/_build/release/presto_cpp/main/presto_server + chmod +x ${GITHUB_WORKSPACE}/presto-native-execution/_build/release/velox/velox/functions/remote/server/velox_functions_remote_server_main + # Ensure transitive dependency libboost-iostreams is found. + ldconfig /usr/local/lib + + - name: Install OpenJDK17 + uses: actions/setup-java@v4 + with: + distribution: 'temurin' + java-version: 17.0.15 + cache: 'maven' + - name: Download nodejs to maven cache + run: .github/bin/download_nodejs + + - name: Maven install + env: + # Use different Maven options to install. + MAVEN_OPTS: "-Xmx2G -XX:+ExitOnOutOfMemoryError" + run: | + for i in $(seq 1 3); do ./mvnw clean install $MAVEN_FAST_INSTALL -pl 'presto-native-tests' -am && s=0 && break || s=$? && sleep 10; done; (exit $s) + + - name: Run presto-on-spark native tests + run: | + export PRESTO_SERVER_PATH="${GITHUB_WORKSPACE}/presto-native-execution/_build/release/presto_cpp/main/presto_server" + export TESTFILES=`find ./presto-native-execution/src/test -type f -name 'TestPrestoSpark*.java'` + # Convert file paths to comma separated class names + export TESTCLASSES=TestPrestoSparkExpressionCompiler,TestPrestoSparkNativeBitwiseFunctionQueries,TestPrestoSparkNativeTpchConnectorQueries,TestPrestoSparkNativeSimpleQueries,TestPrestoSparkSqlFunctions,TestPrestoSparkNativeTpchQueries + for test_file in $TESTFILES + do + tmp=${test_file##*/} + test_class=${tmp%%\.*} + export TESTCLASSES="${TESTCLASSES},$test_class" + done + export TESTCLASSES=${TESTCLASSES#,} + echo "TESTCLASSES = $TESTCLASSES" + + mvn test \ + ${MAVEN_TEST} \ + -pl 'presto-native-execution' \ + -DstorageFormat=${{ matrix.storage-format }} \ + -DsidecarEnabled=${{ matrix.enable-sidecar }} \ + -Dtest="${TESTCLASSES}" \ + -DPRESTO_SERVER=${PRESTO_SERVER_PATH} \ + -DDATA_DIR=${RUNNER_TEMP} \ + -Duser.timezone=America/Bahia_Banderas \ + -T1C + prestocpp-linux-presto-sidecar-tests: needs: [changes, prestocpp-linux-build-for-test] runs-on: ubuntu-22.04 @@ -370,7 +448,7 @@ jobs: # Use different Maven options to install. MAVEN_OPTS: "-Xmx2G -XX:+ExitOnOutOfMemoryError" run: | - for i in $(seq 1 3); do ./mvnw clean install $MAVEN_FAST_INSTALL -pl 'presto-native-execution' -am && s=0 && break || s=$? && sleep 10; done; (exit $s) + for i in $(seq 1 3); do ./mvnw clean install $MAVEN_FAST_INSTALL -pl 'presto-native-sidecar-plugin' -am && s=0 && break || s=$? && sleep 10; done; (exit $s) - name: Run presto-native sidecar tests if: | diff --git a/.github/workflows/tests.yml b/.github/workflows/tests.yml index 0afc4c0fbfa82..3432d50e03e9b 100644 --- a/.github/workflows/tests.yml +++ b/.github/workflows/tests.yml @@ -59,7 +59,7 @@ jobs: - ":presto-hive -P test-hive-parquet" - ":presto-main-base" - ":presto-main" - - ":presto-mongodb -P test-mongo-distributed-queries" + - ":presto-mongodb -P ci-full-tests" - ":presto-redis -P test-redis-integration-smoke-test" - ":presto-elasticsearch" - ":presto-orc" diff --git a/CODEOWNERS b/CODEOWNERS index 60f71c88b99ee..e4a0e3bb56aed 100644 --- a/CODEOWNERS +++ b/CODEOWNERS @@ -50,6 +50,7 @@ /presto-lark-sheets @prestodb/committers /presto-local-file @prestodb/committers /presto-main @prestodb/committers +/presto-main-base @prestodb/committers /presto-matching @prestodb/committers /presto-memory @prestodb/committers /presto-memory-context @prestodb/committers @@ -75,7 +76,9 @@ /presto-resource-group-managers @prestodb/committers /presto-router @prestodb/committers /presto-server @prestodb/committers -/presto-session-property-managers @prestodb/committers +/presto-session-property-managers-common @prestodb/committers +/presto-file-session-property-manager @prestodb/committers +/presto-db-session-property-manager @prestodb/committers /presto-singlestore @prestodb/committers /presto-spi @prestodb/committers /presto-sqlserver @prestodb/committers @@ -100,10 +103,10 @@ CODEOWNERS @prestodb/team-tsc # Presto core # Presto analyzer and optimizer -/presto-main/src/*/java/com/facebook/presto/sql @jaystarshot @feilong-liu @ClarenceThreepwood @prestodb/committers +/presto-main-base/src/*/java/com/facebook/presto/sql @jaystarshot @feilong-liu @ClarenceThreepwood @prestodb/committers # Presto cost based optimizer framework -/presto-main/src/*/java/com/facebook/presto/cost @jaystarshot @feilong-liu @ClarenceThreepwood @prestodb/committers +/presto-main-base/src/*/java/com/facebook/presto/cost @jaystarshot @feilong-liu @ClarenceThreepwood @prestodb/committers # Testing module # Note: all code owners in Presto core should be included here as well @@ -111,14 +114,16 @@ CODEOWNERS @prestodb/team-tsc ##################################################################### # Prestissimo module -/presto-native-execution @prestodb/team-velox -/presto-native-sidecar-plugin @prestodb/team-velox -/presto-native-tests @prestodb/team-velox +/presto-native-execution @prestodb/team-velox @prestodb/committers +/presto-native-sidecar-plugin @pdabre12 @prestodb/team-velox @prestodb/committers +/presto-native-tests @prestodb/team-velox @prestodb/committers +/presto-main-base/src/main/java/com/facebook/presto/sessionpropertyproviders/NativeWorkerSessionPropertyProvider.java @prestodb/team-velox @prestodb/committers /.github/workflows/prestocpp-* @prestodb/team-velox @prestodb/committers ##################################################################### # Presto on Spark module /presto-spark* @shrinidhijoshi @prestodb/committers +/presto-native-execution/*/com/facebook/presto/spark/* @shrinidhijoshi @prestodb/committers ##################################################################### # Presto connectors and plugins @@ -162,4 +167,3 @@ CODEOWNERS @prestodb/team-tsc # Presto CI and builds /.github @czentgr @unidevel @prestodb/committers /docker @czentgr @unidevel @prestodb/committers - diff --git a/README.md b/README.md index 2da36a8a31aea..2b4aa6d998c1d 100644 --- a/README.md +++ b/README.md @@ -59,6 +59,8 @@ Additionally, the Hive plugin must be configured with location of your Hive meta -Dhive.metastore.uri=thrift://localhost:9083 +To modify the loaded plugins in IntelliJ, modify the `config.properties` located in `presto-main/etc`. You can modify `plugin.bundles` with the location of the plugin pom.xml + ### Additional configuration for Java 17 When running with Java 17, additional `--add-opens` flags are required to allow reflective access used by certain catalogs based on which catalogs are configured. diff --git a/pom.xml b/pom.xml index 52ed78bd8ad51..186bcaecac7b0 100644 --- a/pom.xml +++ b/pom.xml @@ -43,14 +43,14 @@ 17 3.3.9 - 4.7.1 + 4.13.2 0.221 ${dep.airlift.version} 0.38 0.6 1.12.782 4.12.0 - 3.4.0 + 3.49.0 19.3.0.0 ${dep.airlift.version} - - ossrh - https://oss.sonatype.org/ - - - com.facebook.presto presto-maven-plugin @@ -3021,7 +3056,7 @@ true ossrh - true + ${release.autoPublish} validated @@ -3139,6 +3174,78 @@ + + spark2 + + + true + + !spark-version + + + + + 2.0.2-6 + + + + presto-spark-classloader-spark2 + + + + + + + org.apache.maven.plugins + maven-enforcer-plugin + 3.3.0 + + + org.codehaus.mojo + extra-enforcer-rules + 1.6.2 + + + + true + + + + + org.codehaus.plexus:plexus-utils + com.google.guava:guava + com.fasterxml.jackson.core:jackson-annotations + com.fasterxml.jackson.core:jackson-core + com.fasterxml.jackson.core:jackson-databind + + + + + + + + org.basepom.maven + duplicate-finder-maven-plugin + + true + + com.github.benmanes.caffeine.* + + META-INF.versions.9.module-info + + META-INF.versions.11.module-info + + META-INF.versions.9.org.apache.lucene.* + + + + + + + + + + spark3 @@ -3150,19 +3257,12 @@ - 3 + 3.4.1-1 - - - - com.facebook.presto.spark - spark-core - 3.4.1-1 - provided - - - + + presto-spark-classloader-spark3 + diff --git a/presto-accumulo/src/main/java/com/facebook/presto/accumulo/AccumuloMetadata.java b/presto-accumulo/src/main/java/com/facebook/presto/accumulo/AccumuloMetadata.java index d585bf8464f88..c1ea2bc36dcb5 100644 --- a/presto-accumulo/src/main/java/com/facebook/presto/accumulo/AccumuloMetadata.java +++ b/presto-accumulo/src/main/java/com/facebook/presto/accumulo/AccumuloMetadata.java @@ -259,7 +259,7 @@ public ConnectorTableHandle getTableHandle(ConnectorSession session, SchemaTable } @Override - public List getTableLayouts( + public ConnectorTableLayoutResult getTableLayoutForConstraint( ConnectorSession session, ConnectorTableHandle table, Constraint constraint, @@ -267,7 +267,7 @@ public List getTableLayouts( { AccumuloTableHandle tableHandle = (AccumuloTableHandle) table; ConnectorTableLayout layout = new ConnectorTableLayout(new AccumuloTableLayoutHandle(tableHandle, constraint.getSummary())); - return ImmutableList.of(new ConnectorTableLayoutResult(layout, constraint.getSummary())); + return new ConnectorTableLayoutResult(layout, constraint.getSummary()); } @Override diff --git a/presto-accumulo/src/test/java/com/facebook/presto/accumulo/TestAccumuloDistributedQueries.java b/presto-accumulo/src/test/java/com/facebook/presto/accumulo/TestAccumuloDistributedQueries.java index a37560fae7a7a..957b10f793dad 100644 --- a/presto-accumulo/src/test/java/com/facebook/presto/accumulo/TestAccumuloDistributedQueries.java +++ b/presto-accumulo/src/test/java/com/facebook/presto/accumulo/TestAccumuloDistributedQueries.java @@ -121,6 +121,12 @@ public void testUpdate() // Updates are not supported by the connector } + @Override + public void testNonAutoCommitTransactionWithRollback() + { + // This connector do not support rollback for insert actions + } + @Override public void testInsert() { diff --git a/presto-analyzer/src/main/java/com/facebook/presto/sql/analyzer/Analysis.java b/presto-analyzer/src/main/java/com/facebook/presto/sql/analyzer/Analysis.java index 66e3f7b0e4dc2..7da15ef222470 100644 --- a/presto-analyzer/src/main/java/com/facebook/presto/sql/analyzer/Analysis.java +++ b/presto-analyzer/src/main/java/com/facebook/presto/sql/analyzer/Analysis.java @@ -14,18 +14,24 @@ package com.facebook.presto.sql.analyzer; import com.facebook.presto.common.QualifiedObjectName; +import com.facebook.presto.common.SourceColumn; import com.facebook.presto.common.Subfield; import com.facebook.presto.common.transaction.TransactionId; import com.facebook.presto.common.type.Type; import com.facebook.presto.spi.ColumnHandle; import com.facebook.presto.spi.ColumnMetadata; +import com.facebook.presto.spi.ConnectorId; import com.facebook.presto.spi.TableHandle; import com.facebook.presto.spi.analyzer.AccessControlInfo; import com.facebook.presto.spi.analyzer.AccessControlInfoForTable; import com.facebook.presto.spi.analyzer.AccessControlReferences; import com.facebook.presto.spi.analyzer.AccessControlRole; +import com.facebook.presto.spi.connector.ConnectorTransactionHandle; +import com.facebook.presto.spi.eventlistener.OutputColumnMetadata; import com.facebook.presto.spi.function.FunctionHandle; import com.facebook.presto.spi.function.FunctionKind; +import com.facebook.presto.spi.function.table.Argument; +import com.facebook.presto.spi.function.table.ConnectorTableFunctionHandle; import com.facebook.presto.spi.security.AccessControl; import com.facebook.presto.spi.security.AccessControlContext; import com.facebook.presto.spi.security.AllowAllAccessControl; @@ -51,6 +57,7 @@ import com.facebook.presto.sql.tree.Statement; import com.facebook.presto.sql.tree.SubqueryExpression; import com.facebook.presto.sql.tree.Table; +import com.facebook.presto.sql.tree.TableFunctionInvocation; import com.google.common.collect.ArrayListMultimap; import com.google.common.collect.HashMultiset; import com.google.common.collect.ImmutableList; @@ -199,6 +206,16 @@ public class Analysis // Keeps track of the subquery we are visiting, so we have access to base query information when processing materialized view status private Optional currentQuerySpecification = Optional.empty(); + // Maps each output Field to its originating SourceColumn(s) for column-level lineage tracking. + private final Multimap originColumnDetails = ArrayListMultimap.create(); + + // Maps each analyzed Expression to the Field(s) it produces, supporting expression-level lineage. + private final Multimap, Field> fieldLineage = ArrayListMultimap.create(); + + private Optional> updatedSourceColumns = Optional.empty(); + + private final Map, TableFunctionInvocationAnalysis> tableFunctionAnalyses = new LinkedHashMap<>(); + public Analysis(@Nullable Statement root, Map, Expression> parameters, boolean isDescribe) { this.root = root; @@ -1039,6 +1056,38 @@ public boolean hasColumnMask(QualifiedObjectName table, String column, String id return columnMaskScopes.contains(new ColumnMaskScopeEntry(table, column, identity)); } + public void addSourceColumns(Field field, Set sourceColumn) + { + originColumnDetails.putAll(field, sourceColumn); + } + + public Set getSourceColumns(Field field) + { + return ImmutableSet.copyOf(originColumnDetails.get(field)); + } + + public void addExpressionFields(Expression expression, Collection fields) + { + fieldLineage.putAll(NodeRef.of(expression), fields); + } + + public Set getExpressionSourceColumns(Expression expression) + { + return fieldLineage.get(NodeRef.of(expression)).stream() + .flatMap(field -> getSourceColumns(field).stream()) + .collect(toImmutableSet()); + } + + public void setUpdatedSourceColumns(Optional> targetColumns) + { + this.updatedSourceColumns = targetColumns; + } + + public Optional> getUpdatedSourceColumns() + { + return updatedSourceColumns; + } + public void registerTableForColumnMasking(QualifiedObjectName table, String column, String identity) { columnMaskScopes.add(new ColumnMaskScopeEntry(table, column, identity)); @@ -1061,6 +1110,16 @@ public Map getColumnMasks(Table table) return columnMasks.getOrDefault(NodeRef.of(table), ImmutableMap.of()); } + public void setTableFunctionAnalysis(TableFunctionInvocation node, TableFunctionInvocationAnalysis analysis) + { + tableFunctionAnalyses.put(NodeRef.of(node), analysis); + } + + public TableFunctionInvocationAnalysis getTableFunctionAnalysis(TableFunctionInvocation node) + { + return tableFunctionAnalyses.get(NodeRef.of(node)); + } + @Immutable public static final class Insert { @@ -1311,4 +1370,57 @@ public int hashCode() return Objects.hash(table, column, identity); } } + + /** + * Encapsulates the result of analyzing a table function invocation. + * Includes the connector ID, function name, argument bindings, and the + * connector-specific table function handle needed for planning and execution. + */ + public static class TableFunctionInvocationAnalysis + { + private final ConnectorId connectorId; + private final String functionName; + private final Map arguments; + private final ConnectorTableFunctionHandle connectorTableFunctionHandle; + private final ConnectorTransactionHandle transactionHandle; + + public TableFunctionInvocationAnalysis( + ConnectorId connectorId, + String functionName, + Map arguments, + ConnectorTableFunctionHandle connectorTableFunctionHandle, + ConnectorTransactionHandle transactionHandle) + { + this.connectorId = requireNonNull(connectorId, "connectorId is null"); + this.functionName = requireNonNull(functionName, "functionName is null"); + this.arguments = ImmutableMap.copyOf(arguments); + this.connectorTableFunctionHandle = requireNonNull(connectorTableFunctionHandle, "connectorTableFunctionHandle is null"); + this.transactionHandle = requireNonNull(transactionHandle, "transactionHandle is null"); + } + + public ConnectorId getConnectorId() + { + return connectorId; + } + + public String getFunctionName() + { + return functionName; + } + + public Map getArguments() + { + return arguments; + } + + public ConnectorTableFunctionHandle getConnectorTableFunctionHandle() + { + return connectorTableFunctionHandle; + } + + public ConnectorTransactionHandle getTransactionHandle() + { + return transactionHandle; + } + } } diff --git a/presto-analyzer/src/main/java/com/facebook/presto/sql/analyzer/Field.java b/presto-analyzer/src/main/java/com/facebook/presto/sql/analyzer/Field.java index 630f4670f6cc2..15c33950ce14b 100644 --- a/presto-analyzer/src/main/java/com/facebook/presto/sql/analyzer/Field.java +++ b/presto-analyzer/src/main/java/com/facebook/presto/sql/analyzer/Field.java @@ -86,6 +86,14 @@ public Field(Optional nodeLocation, Optional relati this.aliased = aliased; } + public static Field newUnqualified(Optional name, Type type) + { + requireNonNull(name, "name is null"); + requireNonNull(type, "type is null"); + + return new Field(Optional.empty(), Optional.empty(), name, type, false, Optional.empty(), Optional.empty(), false); + } + public Optional getNodeLocation() { return nodeLocation; diff --git a/presto-analyzer/src/main/java/com/facebook/presto/sql/analyzer/SemanticErrorCode.java b/presto-analyzer/src/main/java/com/facebook/presto/sql/analyzer/SemanticErrorCode.java index 656b6dbaab079..5c9be03540f67 100644 --- a/presto-analyzer/src/main/java/com/facebook/presto/sql/analyzer/SemanticErrorCode.java +++ b/presto-analyzer/src/main/java/com/facebook/presto/sql/analyzer/SemanticErrorCode.java @@ -58,6 +58,13 @@ public enum SemanticErrorCode DUPLICATE_PARAMETER_NAME, EXCEPTIONS_WHEN_RESOLVING_FUNCTIONS, + MISSING_RETURN_TYPE, + AMBIGUOUS_RETURN_TYPE, + MISSING_ARGUMENT, + INVALID_FUNCTION_ARGUMENT, + INVALID_ARGUMENTS, + NOT_IMPLEMENTED, + ORDER_BY_MUST_BE_IN_SELECT, ORDER_BY_MUST_BE_IN_AGGREGATE, REFERENCE_TO_OUTPUT_ATTRIBUTE_WITHIN_ORDER_BY_GROUPING, diff --git a/presto-atop/src/main/java/com/facebook/presto/atop/AtopMetadata.java b/presto-atop/src/main/java/com/facebook/presto/atop/AtopMetadata.java index 04f2e84478935..16710c146f604 100644 --- a/presto-atop/src/main/java/com/facebook/presto/atop/AtopMetadata.java +++ b/presto-atop/src/main/java/com/facebook/presto/atop/AtopMetadata.java @@ -85,7 +85,8 @@ public ConnectorTableHandle getTableHandle(ConnectorSession session, SchemaTable } @Override - public List getTableLayouts(ConnectorSession session, + public ConnectorTableLayoutResult getTableLayoutForConstraint( + ConnectorSession session, ConnectorTableHandle table, Constraint constraint, Optional> desiredColumns) @@ -104,7 +105,7 @@ public List getTableLayouts(ConnectorSession session } AtopTableLayoutHandle layoutHandle = new AtopTableLayoutHandle(tableHandle, startTimeDomain, endTimeDomain); ConnectorTableLayout tableLayout = getTableLayout(session, layoutHandle); - return ImmutableList.of(new ConnectorTableLayoutResult(tableLayout, constraint.getSummary())); + return new ConnectorTableLayoutResult(tableLayout, constraint.getSummary()); } @Override diff --git a/presto-base-arrow-flight/src/main/java/com/facebook/plugin/arrow/ArrowMetadata.java b/presto-base-arrow-flight/src/main/java/com/facebook/plugin/arrow/ArrowMetadata.java index 4018e8f25fbaf..15dda40404686 100644 --- a/presto-base-arrow-flight/src/main/java/com/facebook/plugin/arrow/ArrowMetadata.java +++ b/presto-base-arrow-flight/src/main/java/com/facebook/plugin/arrow/ArrowMetadata.java @@ -112,7 +112,11 @@ public Map getColumnHandles(ConnectorSession session, Conn } @Override - public List getTableLayouts(ConnectorSession session, ConnectorTableHandle table, Constraint constraint, Optional> desiredColumns) + public ConnectorTableLayoutResult getTableLayoutForConstraint( + ConnectorSession session, + ConnectorTableHandle table, + Constraint constraint, + Optional> desiredColumns) { checkArgument(table instanceof ArrowTableHandle, "Invalid table handle: Expected an instance of ArrowTableHandle but received %s", @@ -129,7 +133,7 @@ public List getTableLayouts(ConnectorSession session } ConnectorTableLayout layout = new ConnectorTableLayout(new ArrowTableLayoutHandle(tableHandle, columns, constraint.getSummary())); - return ImmutableList.of(new ConnectorTableLayoutResult(layout, constraint.getSummary())); + return new ConnectorTableLayoutResult(layout, constraint.getSummary()); } @Override diff --git a/presto-base-arrow-flight/src/test/java/com/facebook/plugin/arrow/ArrowFlightQueryRunner.java b/presto-base-arrow-flight/src/test/java/com/facebook/plugin/arrow/ArrowFlightQueryRunner.java index eea7ff598da72..66db074b419b8 100644 --- a/presto-base-arrow-flight/src/test/java/com/facebook/plugin/arrow/ArrowFlightQueryRunner.java +++ b/presto-base-arrow-flight/src/test/java/com/facebook/plugin/arrow/ArrowFlightQueryRunner.java @@ -39,6 +39,8 @@ public class ArrowFlightQueryRunner { + private static final Logger log = Logger.get(ArrowFlightQueryRunner.class); + private ArrowFlightQueryRunner() { throw new UnsupportedOperationException("This is a utility class and cannot be instantiated"); @@ -54,24 +56,15 @@ public static int findUnusedPort() public static DistributedQueryRunner createQueryRunner(int flightServerPort) throws Exception { - return createQueryRunner(flightServerPort, ImmutableMap.of(), ImmutableMap.of(), Optional.empty()); + return createQueryRunner(flightServerPort, ImmutableMap.of(), ImmutableMap.of(), Optional.empty(), Optional.empty()); } public static DistributedQueryRunner createQueryRunner( int flightServerPort, Map extraProperties, Map coordinatorProperties, - Optional> externalWorkerLauncher) - throws Exception - { - return createQueryRunner(extraProperties, ImmutableMap.of("arrow-flight.server.port", String.valueOf(flightServerPort)), coordinatorProperties, externalWorkerLauncher); - } - - private static DistributedQueryRunner createQueryRunner( - Map extraProperties, - Map catalogProperties, - Map coordinatorProperties, - Optional> externalWorkerLauncher) + Optional> externalWorkerLauncher, + Optional mTLSEnabled) throws Exception { Session session = testSessionBuilder() @@ -86,17 +79,24 @@ private static DistributedQueryRunner createQueryRunner( DistributedQueryRunner queryRunner = queryRunnerBuilder .setExtraProperties(extraProperties) .setCoordinatorProperties(coordinatorProperties) - .setExternalWorkerLauncher(externalWorkerLauncher).build(); + .setExternalWorkerLauncher(externalWorkerLauncher) + .build(); try { boolean nativeExecution = externalWorkerLauncher.isPresent(); queryRunner.installPlugin(new TestingArrowFlightPlugin(nativeExecution)); + Map catalogProperties = ImmutableMap.of("arrow-flight.server.port", String.valueOf(flightServerPort)); ImmutableMap.Builder properties = ImmutableMap.builder() .putAll(catalogProperties) .put("arrow-flight.server", "localhost") .put("arrow-flight.server-ssl-enabled", "true") - .put("arrow-flight.server-ssl-certificate", "src/test/resources/server.crt"); + .put("arrow-flight.server-ssl-certificate", "src/test/resources/certs/ca.crt"); + + if (mTLSEnabled.orElse(false)) { + properties.put("arrow-flight.client-ssl-certificate", "src/test/resources/certs/client.crt"); + properties.put("arrow-flight.client-ssl-key", "src/test/resources/certs/client.key"); + } queryRunner.createCatalog(ARROW_FLIGHT_CATALOG, ARROW_FLIGHT_CONNECTOR, properties.build()); @@ -125,25 +125,33 @@ public static void main(String[] args) { Logging.initialize(); + boolean mTLSenabled = Boolean.parseBoolean(System.getProperty("flight.mtls.enabled", "false")); + RootAllocator allocator = new RootAllocator(Long.MAX_VALUE); Location serverLocation = Location.forGrpcTls("localhost", 9443); - File certChainFile = new File("src/test/resources/server.crt"); - File privateKeyFile = new File("src/test/resources/server.key"); - FlightServer server = FlightServer.builder(allocator, serverLocation, new TestingArrowProducer(allocator)) - .useTls(certChainFile, privateKeyFile) - .build(); + FlightServer.Builder serverBuilder = FlightServer.builder(allocator, serverLocation, new TestingArrowProducer(allocator)); - server.start(); + File serverCert = new File("src/test/resources/certs/server.crt"); + File serverKey = new File("src/test/resources/certs/server.key"); + serverBuilder.useTls(serverCert, serverKey); - Logger log = Logger.get(ArrowFlightQueryRunner.class); + if (mTLSenabled) { + File caCert = new File("src/test/resources/certs/ca.crt"); + serverBuilder.useMTlsClientVerification(caCert); + } + + FlightServer server = serverBuilder.build(); + server.start(); log.info("Server listening on port " + server.getPort()); DistributedQueryRunner queryRunner = createQueryRunner( + server.getPort(), ImmutableMap.of("http-server.http.port", "8080"), - ImmutableMap.of("arrow-flight.server.port", String.valueOf(9443)), ImmutableMap.of(), - Optional.empty()); + Optional.empty(), + Optional.of(mTLSenabled)); + Thread.sleep(10); log.info("======== SERVER STARTED ========"); log.info("\n====\n%s\n====", queryRunner.getCoordinator().getBaseUrl()); diff --git a/presto-base-arrow-flight/src/test/java/com/facebook/plugin/arrow/TestArrowFlightEchoQueries.java b/presto-base-arrow-flight/src/test/java/com/facebook/plugin/arrow/TestArrowFlightEchoQueries.java index 8715d413f687f..f658fb6f805c1 100644 --- a/presto-base-arrow-flight/src/test/java/com/facebook/plugin/arrow/TestArrowFlightEchoQueries.java +++ b/presto-base-arrow-flight/src/test/java/com/facebook/plugin/arrow/TestArrowFlightEchoQueries.java @@ -131,8 +131,8 @@ public void setup() throws Exception { arrowFlightQueryRunner = getDistributedQueryRunner(); - File certChainFile = new File("src/test/resources/server.crt"); - File privateKeyFile = new File("src/test/resources/server.key"); + File certChainFile = new File("src/test/resources/certs/server.crt"); + File privateKeyFile = new File("src/test/resources/certs/server.key"); allocator = new RootAllocator(Long.MAX_VALUE); @@ -407,7 +407,7 @@ private static MapType createMapType(Type keyType, Type valueType) private static FlightClient createFlightClient(BufferAllocator allocator, int serverPort) throws IOException { - InputStream trustedCertificate = new ByteArrayInputStream(Files.readAllBytes(Paths.get("src/test/resources/server.crt"))); + InputStream trustedCertificate = new ByteArrayInputStream(Files.readAllBytes(Paths.get("src/test/resources/certs/server.crt"))); Location location = Location.forGrpcTls("localhost", serverPort); return FlightClient.builder(allocator, location).useTls().trustedCertificates(trustedCertificate).build(); } diff --git a/presto-base-arrow-flight/src/test/java/com/facebook/plugin/arrow/TestArrowFlightIntegrationSmokeTest.java b/presto-base-arrow-flight/src/test/java/com/facebook/plugin/arrow/TestArrowFlightIntegrationSmokeTest.java index 1bca76c23a9f0..916dff5e964e2 100644 --- a/presto-base-arrow-flight/src/test/java/com/facebook/plugin/arrow/TestArrowFlightIntegrationSmokeTest.java +++ b/presto-base-arrow-flight/src/test/java/com/facebook/plugin/arrow/TestArrowFlightIntegrationSmokeTest.java @@ -48,8 +48,8 @@ public void setup() throws Exception { arrowFlightQueryRunner = getDistributedQueryRunner(); - File certChainFile = new File("src/test/resources/server.crt"); - File privateKeyFile = new File("src/test/resources/server.key"); + File certChainFile = new File("src/test/resources/certs/server.crt"); + File privateKeyFile = new File("src/test/resources/certs/server.key"); allocator = new RootAllocator(Long.MAX_VALUE); Location location = Location.forGrpcTls("127.0.0.1", serverPort); diff --git a/presto-base-arrow-flight/src/test/java/com/facebook/plugin/arrow/TestArrowFlightMtls.java b/presto-base-arrow-flight/src/test/java/com/facebook/plugin/arrow/TestArrowFlightMtls.java index 169f08db4ddb6..f0f3b8aa76ed0 100644 --- a/presto-base-arrow-flight/src/test/java/com/facebook/plugin/arrow/TestArrowFlightMtls.java +++ b/presto-base-arrow-flight/src/test/java/com/facebook/plugin/arrow/TestArrowFlightMtls.java @@ -60,9 +60,9 @@ private void setup() arrowFlightQueryRunner.createCatalog(ARROW_FLIGHT_CATALOG_WITH_NO_MTLS_CERTS, ARROW_FLIGHT_CONNECTOR, getNoMtlsCatalogProperties()); arrowFlightQueryRunner.createCatalog(ARROW_FLIGHT_CATALOG_WITH_MTLS_CERTS, ARROW_FLIGHT_CONNECTOR, getMtlsCatalogProperties()); - File certChainFile = new File("src/test/resources/mtls/server.crt"); - File privateKeyFile = new File("src/test/resources/mtls/server.key"); - File caCertFile = new File("src/test/resources/mtls/ca.crt"); + File certChainFile = new File("src/test/resources/certs/server.crt"); + File privateKeyFile = new File("src/test/resources/certs/server.key"); + File caCertFile = new File("src/test/resources/certs/ca.crt"); allocator = new RootAllocator(Long.MAX_VALUE); @@ -89,7 +89,7 @@ private void tearDown() protected QueryRunner createQueryRunner() throws Exception { - return ArrowFlightQueryRunner.createQueryRunner(serverPort, ImmutableMap.of(), ImmutableMap.of(), Optional.empty()); + return ArrowFlightQueryRunner.createQueryRunner(serverPort, ImmutableMap.of(), ImmutableMap.of(), Optional.empty(), Optional.empty()); } private Map getInvalidCertCatalogProperties() @@ -98,9 +98,9 @@ private Map getInvalidCertCatalogProperties() .put("arrow-flight.server.port", String.valueOf(serverPort)) .put("arrow-flight.server", "localhost") .put("arrow-flight.server-ssl-enabled", "true") - .put("arrow-flight.server-ssl-certificate", "src/test/resources/mtls/server.crt") - .put("arrow-flight.client-ssl-certificate", "src/test/resources/mtls/invalid_cert.crt") - .put("arrow-flight.client-ssl-key", "src/test/resources/mtls/client.key"); + .put("arrow-flight.server-ssl-certificate", "src/test/resources/certs/server.crt") + .put("arrow-flight.client-ssl-certificate", "src/test/resources/certs/invalid_cert.crt") + .put("arrow-flight.client-ssl-key", "src/test/resources/certs/client.key"); return catalogProperties.build(); } @@ -110,7 +110,7 @@ private Map getNoMtlsCatalogProperties() .put("arrow-flight.server.port", String.valueOf(serverPort)) .put("arrow-flight.server", "localhost") .put("arrow-flight.server-ssl-enabled", "true") - .put("arrow-flight.server-ssl-certificate", "src/test/resources/mtls/server.crt"); + .put("arrow-flight.server-ssl-certificate", "src/test/resources/certs/server.crt"); return catalogProperties.build(); } @@ -120,9 +120,9 @@ private Map getMtlsCatalogProperties() .put("arrow-flight.server.port", String.valueOf(serverPort)) .put("arrow-flight.server", "localhost") .put("arrow-flight.server-ssl-enabled", "true") - .put("arrow-flight.server-ssl-certificate", "src/test/resources/mtls/server.crt") - .put("arrow-flight.client-ssl-certificate", "src/test/resources/mtls/client.crt") - .put("arrow-flight.client-ssl-key", "src/test/resources/mtls/client.key"); + .put("arrow-flight.server-ssl-certificate", "src/test/resources/certs/server.crt") + .put("arrow-flight.client-ssl-certificate", "src/test/resources/certs/client.crt") + .put("arrow-flight.client-ssl-key", "src/test/resources/certs/client.key"); return catalogProperties.build(); } diff --git a/presto-base-arrow-flight/src/test/java/com/facebook/plugin/arrow/TestArrowFlightNativeQueries.java b/presto-base-arrow-flight/src/test/java/com/facebook/plugin/arrow/TestArrowFlightNativeQueries.java index 21ee7f53ed78a..d6f945c45853c 100644 --- a/presto-base-arrow-flight/src/test/java/com/facebook/plugin/arrow/TestArrowFlightNativeQueries.java +++ b/presto-base-arrow-flight/src/test/java/com/facebook/plugin/arrow/TestArrowFlightNativeQueries.java @@ -60,23 +60,32 @@ public TestArrowFlightNativeQueries() this.serverPort = ArrowFlightQueryRunner.findUnusedPort(); } + protected boolean ismTLSEnabled() + { + return false; + } + @BeforeClass public void setup() throws Exception { arrowFlightQueryRunner = getDistributedQueryRunner(); - allocator = new RootAllocator(Long.MAX_VALUE); Location location = Location.forGrpcTls("localhost", serverPort); - File certChainFile = new File("src/test/resources/server.crt"); - File privateKeyFile = new File("src/test/resources/server.key"); + FlightServer.Builder serverBuilder = FlightServer.builder(allocator, location, new TestingArrowProducer(allocator)); - server = FlightServer.builder(allocator, location, new TestingArrowProducer(allocator)) - .useTls(certChainFile, privateKeyFile) - .build(); + File serverCert = new File("src/test/resources/certs/server.crt"); + File serverKey = new File("src/test/resources/certs/server.key"); + serverBuilder.useTls(serverCert, serverKey); + + if (ismTLSEnabled()) { + File caCert = new File("src/test/resources/certs/ca.crt"); + serverBuilder.useMTlsClientVerification(caCert); + } + server = serverBuilder.build(); server.start(); - log.info("Server listening on port %s", server.getPort()); + log.info("Server listening on port %s (%s)", server.getPort(), ismTLSEnabled() ? "mTLS" : "TLS"); } @AfterClass(alwaysRun = true) @@ -99,9 +108,13 @@ protected QueryRunner createQueryRunner() log.info("Using PRESTO_SERVER binary at %s", prestoServerPath); ImmutableMap coordinatorProperties = ImmutableMap.of("native-execution-enabled", "true"); - String flightCertPath = Paths.get("src/test/resources/server.crt").toAbsolutePath().toString(); - return ArrowFlightQueryRunner.createQueryRunner(serverPort, getNativeWorkerSystemProperties(), coordinatorProperties, getExternalWorkerLauncher(prestoServerPath.toString(), serverPort, flightCertPath)); + return ArrowFlightQueryRunner.createQueryRunner( + serverPort, + getNativeWorkerSystemProperties(), + coordinatorProperties, + getExternalWorkerLauncher(prestoServerPath.toString(), serverPort, ismTLSEnabled()), + Optional.of(ismTLSEnabled())); } @Override @@ -338,50 +351,63 @@ public static Map getNativeWorkerSystemProperties() .build(); } - public static Optional> getExternalWorkerLauncher(String prestoServerPath, int flightServerPort, String flightCertPath) + public static Optional> getExternalWorkerLauncher(String prestoServerPath, int flightServerPort, boolean ismTLSEnabled) { - return - Optional.of((workerIndex, discoveryUri) -> { - try { - Path dir = Paths.get("/tmp", TestArrowFlightNativeQueries.class.getSimpleName()); - Files.createDirectories(dir); - Path tempDirectoryPath = Files.createTempDirectory(dir, "worker"); - log.info("Temp directory for Worker #%d: %s", workerIndex, tempDirectoryPath.toString()); - - // Write config file - use an ephemeral port for the worker. - String configProperties = format("discovery.uri=%s%n" + - "presto.version=testversion%n" + - "system-memory-gb=4%n" + - "http-server.http.port=0%n", discoveryUri); - - Files.write(tempDirectoryPath.resolve("config.properties"), configProperties.getBytes()); - Files.write(tempDirectoryPath.resolve("node.properties"), - format("node.id=%s%n" + - "node.internal-address=127.0.0.1%n" + - "node.environment=testing%n" + - "node.location=test-location", UUID.randomUUID()).getBytes()); - - Path catalogDirectoryPath = tempDirectoryPath.resolve("catalog"); - Files.createDirectory(catalogDirectoryPath); - - Files.write(catalogDirectoryPath.resolve(format("%s.properties", ARROW_FLIGHT_CATALOG)), - format("connector.name=%s\n" + - "arrow-flight.server=localhost\n" + - "arrow-flight.server.port=%d\n" + - "arrow-flight.server-ssl-enabled=true\n" + - "arrow-flight.server-ssl-certificate=%s", ARROW_FLIGHT_CONNECTOR, flightServerPort, flightCertPath).getBytes()); - - // Disable stack trace capturing as some queries (using TRY) generate a lot of exceptions. - return new ProcessBuilder(prestoServerPath, "--logtostderr=1", "--v=1") - .directory(tempDirectoryPath.toFile()) - .redirectErrorStream(true) - .redirectOutput(ProcessBuilder.Redirect.to(tempDirectoryPath.resolve("worker." + workerIndex + ".out").toFile())) - .redirectError(ProcessBuilder.Redirect.to(tempDirectoryPath.resolve("worker." + workerIndex + ".out").toFile())) - .start(); - } - catch (IOException e) { - throw new UncheckedIOException(e); - } - }); + return Optional.of((workerIndex, discoveryUri) -> { + try { + Path dir = Paths.get("/tmp", TestArrowFlightNativeQueries.class.getSimpleName()); + Files.createDirectories(dir); + Path tempDirectoryPath = Files.createTempDirectory(dir, "worker"); + log.info("Temp directory for Worker #%d: %s", workerIndex, tempDirectoryPath.toString()); + + // Write config file - use an ephemeral port for the worker. + String configProperties = format("discovery.uri=%s%n" + + "presto.version=testversion%n" + + "system-memory-gb=4%n" + + "http-server.http.port=0%n", discoveryUri); + + Files.write(tempDirectoryPath.resolve("config.properties"), configProperties.getBytes()); + Files.write(tempDirectoryPath.resolve("node.properties"), + format("node.id=%s%n" + + "node.internal-address=127.0.0.1%n" + + "node.environment=testing%n" + + "node.location=test-location", UUID.randomUUID()).getBytes()); + + Path catalogDirectoryPath = tempDirectoryPath.resolve("catalog"); + Files.createDirectory(catalogDirectoryPath); + + String caCertPath = Paths.get("src/test/resources/certs/ca.crt").toAbsolutePath().toString(); + + StringBuilder catalogBuilder = new StringBuilder(); + catalogBuilder.append(format( + "connector.name=%s\n" + + "arrow-flight.server=localhost\n" + + "arrow-flight.server.port=%d\n" + + "arrow-flight.server-ssl-enabled=true\n" + + "arrow-flight.server-ssl-certificate=%s\n", + ARROW_FLIGHT_CONNECTOR, flightServerPort, caCertPath)); + + if (ismTLSEnabled) { + String clientCertPath = Paths.get("src/test/resources/certs/client.crt").toAbsolutePath().toString(); + String clientKeyPath = Paths.get("src/test/resources/certs/client.key").toAbsolutePath().toString(); + catalogBuilder.append(format("arrow-flight.client-ssl-certificate=%s\n", clientCertPath)); + catalogBuilder.append(format("arrow-flight.client-ssl-key=%s\n", clientKeyPath)); + } + + Files.write( + catalogDirectoryPath.resolve(format("%s.properties", ARROW_FLIGHT_CATALOG)), + catalogBuilder.toString().getBytes()); + + return new ProcessBuilder(prestoServerPath, "--logtostderr=1", "--v=1") + .directory(tempDirectoryPath.toFile()) + .redirectErrorStream(true) + .redirectOutput(ProcessBuilder.Redirect.to(tempDirectoryPath.resolve("worker." + workerIndex + ".out").toFile())) + .redirectError(ProcessBuilder.Redirect.to(tempDirectoryPath.resolve("worker." + workerIndex + ".out").toFile())) + .start(); + } + catch (IOException e) { + throw new UncheckedIOException(e); + } + }); } } diff --git a/presto-base-arrow-flight/src/test/java/com/facebook/plugin/arrow/TestArrowFlightNativeQueriesMtls.java b/presto-base-arrow-flight/src/test/java/com/facebook/plugin/arrow/TestArrowFlightNativeQueriesMtls.java new file mode 100644 index 0000000000000..25079773e7c02 --- /dev/null +++ b/presto-base-arrow-flight/src/test/java/com/facebook/plugin/arrow/TestArrowFlightNativeQueriesMtls.java @@ -0,0 +1,33 @@ +/* + * Licensed 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 com.facebook.plugin.arrow; + +import java.io.IOException; + +public class TestArrowFlightNativeQueriesMtls + extends TestArrowFlightNativeQueries +{ + public TestArrowFlightNativeQueriesMtls() + throws IOException + { + super(); + } + + @Override + protected boolean ismTLSEnabled() + { + return true; + } +} diff --git a/presto-base-arrow-flight/src/test/java/com/facebook/plugin/arrow/TestArrowFlightQueries.java b/presto-base-arrow-flight/src/test/java/com/facebook/plugin/arrow/TestArrowFlightQueries.java index 9be967eab2c08..cc1edc0735ce0 100644 --- a/presto-base-arrow-flight/src/test/java/com/facebook/plugin/arrow/TestArrowFlightQueries.java +++ b/presto-base-arrow-flight/src/test/java/com/facebook/plugin/arrow/TestArrowFlightQueries.java @@ -70,8 +70,8 @@ public void setup() throws Exception { arrowFlightQueryRunner = getDistributedQueryRunner(); - File certChainFile = new File("src/test/resources/server.crt"); - File privateKeyFile = new File("src/test/resources/server.key"); + File certChainFile = new File("src/test/resources/certs/server.crt"); + File privateKeyFile = new File("src/test/resources/certs/server.key"); allocator = new RootAllocator(Long.MAX_VALUE); Location location = Location.forGrpcTls("localhost", serverPort); diff --git a/presto-base-arrow-flight/src/test/resources/mtls/ca.crt b/presto-base-arrow-flight/src/test/resources/certs/ca.crt similarity index 100% rename from presto-base-arrow-flight/src/test/resources/mtls/ca.crt rename to presto-base-arrow-flight/src/test/resources/certs/ca.crt diff --git a/presto-base-arrow-flight/src/test/resources/mtls/client.crt b/presto-base-arrow-flight/src/test/resources/certs/client.crt similarity index 100% rename from presto-base-arrow-flight/src/test/resources/mtls/client.crt rename to presto-base-arrow-flight/src/test/resources/certs/client.crt diff --git a/presto-base-arrow-flight/src/test/resources/mtls/client.key b/presto-base-arrow-flight/src/test/resources/certs/client.key similarity index 100% rename from presto-base-arrow-flight/src/test/resources/mtls/client.key rename to presto-base-arrow-flight/src/test/resources/certs/client.key diff --git a/presto-base-arrow-flight/src/test/resources/mtls/invalid_cert.crt b/presto-base-arrow-flight/src/test/resources/certs/invalid_cert.crt similarity index 100% rename from presto-base-arrow-flight/src/test/resources/mtls/invalid_cert.crt rename to presto-base-arrow-flight/src/test/resources/certs/invalid_cert.crt diff --git a/presto-base-arrow-flight/src/test/resources/mtls/server.crt b/presto-base-arrow-flight/src/test/resources/certs/server.crt similarity index 100% rename from presto-base-arrow-flight/src/test/resources/mtls/server.crt rename to presto-base-arrow-flight/src/test/resources/certs/server.crt diff --git a/presto-base-arrow-flight/src/test/resources/mtls/server.key b/presto-base-arrow-flight/src/test/resources/certs/server.key similarity index 100% rename from presto-base-arrow-flight/src/test/resources/mtls/server.key rename to presto-base-arrow-flight/src/test/resources/certs/server.key diff --git a/presto-base-arrow-flight/src/test/resources/server.crt b/presto-base-arrow-flight/src/test/resources/server.crt deleted file mode 100644 index f070253c9e119..0000000000000 --- a/presto-base-arrow-flight/src/test/resources/server.crt +++ /dev/null @@ -1,32 +0,0 @@ ------BEGIN CERTIFICATE----- -MIIFlTCCA32gAwIBAgIUf2p3qdxzpsOofYDpXNXDJA1fjbgwDQYJKoZIhvcNAQEL -BQAwWTELMAkGA1UEBhMCQVUxEzARBgNVBAgMClNvbWUtU3RhdGUxITAfBgNVBAoM -GEludGVybmV0IFdpZGdpdHMgUHR5IEx0ZDESMBAGA1UEAwwJbG9jYWxob3N0MCAX -DTI0MTIyMTA3MzUwMFoYDzIxMjQxMTI3MDczNTAwWjBZMQswCQYDVQQGEwJBVTET -MBEGA1UECAwKU29tZS1TdGF0ZTEhMB8GA1UECgwYSW50ZXJuZXQgV2lkZ2l0cyBQ -dHkgTHRkMRIwEAYDVQQDDAlsb2NhbGhvc3QwggIiMA0GCSqGSIb3DQEBAQUAA4IC -DwAwggIKAoICAQCmsVs+OHGQ/GlR9+KpgcKkBvXZQAUtCJzFG27GXBYoltJDAxsi -o47KbGkUWxf5E+4oFbbz7s8OHSuQnjP2XVHbKvmgicf3uP98fKkQtjvl+xn8vT18 -s6a8Kv8hj+f6MRWGwpHa/sQKU9uZmRYTRh+32vDZRtAvmpzkf6+2K8B1fFbwVmuC -j3ULb+0iYHnomC3aMWBFXkxjEmsamx4YK74NtQU98+EjQZwhWgWXhW5shS1kSs4r -3N6++tonBz+tDKAhCMueRRJAQXGjKqL7qDZn7wpk53L/fZT9mgRYyA+PN2ND9L0H -nMGMjJl71p42kkgIGOpllsmK6+g0Bj6aC/uCEnX2AtM0g57Th7U2aLwgOmRshC0s -uuBHxMWUgzJsB1dscXrFPPB+XVcUlgcwRbGsQG5VzK/rYRV4y1FmF5vSLY60mF43 -hFDAcnh4mnnBkobca5Dl6PSUpFmDkF56IUgYQCTrE6hPYnrIJKhPcfpmr/tm/Acr -ra1sPp/QPSFIxI9j7Nzm/QsOBF3Zy4AbbbOmhJOjNtwEi59r9za/FhxV5kSN/YM9 -HyyYYebxW/jXF/7hzQzWYfJBz1SgdD9prl4ml8VMVJdhZmBTzhVciPXizRUKkbD+ -LvQKw8q4a24/VruUvW15J39qalhdyWf3vqGuORWowpG7oYnGXik3kYVT5QIDAQAB -o1MwUTAdBgNVHQ4EFgQUkkKtG5568IDoFtn3AK0Yes0CqGgwHwYDVR0jBBgwFoAU -kkKtG5568IDoFtn3AK0Yes0CqGgwDwYDVR0TAQH/BAUwAwEB/zANBgkqhkiG9w0B -AQsFAAOCAgEAVjwxyO5N4EGQHf9euKiTjksUMeJrJkPMFYS/1UmQB5/yc0TkOxv3 -M8yvTzkzxWAQV2X+CPbOLgqO2eCoSev2GHVUyHLLUEFATzAgTpF79Ovzq6JkSgFo -GFxTL0TuA+d++430wYyts7iYY9uT5Xt8xp4+SwXrCvLWxHosP3xGUmNvY4sP0hdS -beIvdGE/J8izgy5DRt2fWZ03mmwfKqiz/qhKGj9DDsHkA/1jyKIivP/nufr9dzDr -41lhk1N7qFWkOjbMd06NYySIe0MaapIkenjT1IOgqGw2f98RfSEomoaXuxN0VoSM -6dZ4rN97cER25X7/zE0zCZurjCLHzPTyuTspYGEK+9U4plOeWh0keQEcdpcHTAr+ -NqU3VlhXVxz91nVREpRJmKk+r4c+xdrfY3YkDcSJ1dazbd3eS05Ggx51KOqes8Zb -hFQfhIDqvaqXDNlkBezLpr4v/MU69+cp7SOn5uPnOccS6sd7fzl/PUZhEjd5ZIws -8SX79OwhQjbYZYRHJSPPasb8B1amULtoo0pJ5izSkXxileiGuXhRO5stiBux7SL+ -oJAztjuRf0IvP6LWOMHgqquzc2JiEDCz0DPnTCqoXGZlT2HPGzXOoDSTOmRdFx+L -qi/DeY+MpIMVov/rplqjydXw6AuQDxcV1GvyjMvaHxJG5MEBC/mVeqQ= ------END CERTIFICATE----- \ No newline at end of file diff --git a/presto-base-arrow-flight/src/test/resources/server.key b/presto-base-arrow-flight/src/test/resources/server.key deleted file mode 100644 index 4524c26943f2d..0000000000000 --- a/presto-base-arrow-flight/src/test/resources/server.key +++ /dev/null @@ -1,52 +0,0 @@ ------BEGIN PRIVATE KEY----- -MIIJQwIBADANBgkqhkiG9w0BAQEFAASCCS0wggkpAgEAAoICAQCmsVs+OHGQ/GlR -9+KpgcKkBvXZQAUtCJzFG27GXBYoltJDAxsio47KbGkUWxf5E+4oFbbz7s8OHSuQ -njP2XVHbKvmgicf3uP98fKkQtjvl+xn8vT18s6a8Kv8hj+f6MRWGwpHa/sQKU9uZ -mRYTRh+32vDZRtAvmpzkf6+2K8B1fFbwVmuCj3ULb+0iYHnomC3aMWBFXkxjEmsa -mx4YK74NtQU98+EjQZwhWgWXhW5shS1kSs4r3N6++tonBz+tDKAhCMueRRJAQXGj -KqL7qDZn7wpk53L/fZT9mgRYyA+PN2ND9L0HnMGMjJl71p42kkgIGOpllsmK6+g0 -Bj6aC/uCEnX2AtM0g57Th7U2aLwgOmRshC0suuBHxMWUgzJsB1dscXrFPPB+XVcU -lgcwRbGsQG5VzK/rYRV4y1FmF5vSLY60mF43hFDAcnh4mnnBkobca5Dl6PSUpFmD -kF56IUgYQCTrE6hPYnrIJKhPcfpmr/tm/Acrra1sPp/QPSFIxI9j7Nzm/QsOBF3Z -y4AbbbOmhJOjNtwEi59r9za/FhxV5kSN/YM9HyyYYebxW/jXF/7hzQzWYfJBz1Sg -dD9prl4ml8VMVJdhZmBTzhVciPXizRUKkbD+LvQKw8q4a24/VruUvW15J39qalhd -yWf3vqGuORWowpG7oYnGXik3kYVT5QIDAQABAoICAE0MeIzTgSbPjQz+w82u9V1k -/DlNfrb4nqH7EqJsSS+8uvaPlnjV2fgV0SJAEt4mCLSNiPHKpfkzoYHopkMPknj4 -Lcc3OG94GtubMXhQi3I7tSDeBfBAh+a9Bw2n20WJb5ZJFCsCDHJrnXsrSAljpeCR -OjdsJGmEkVWK8ZiGM6D6dqMDhxEjpynAs/7qUh8hTDxpC0M1GaDHkCsNnQT2HxRt -4jznH97wgi7mUeReIBLYIgmUDCU5I9ppz/EvSA8AYXmze46uBYge19xgJlKlR3SW -CJtoYf7XOMlZ6f1xh8OeiesM0l0U51/EU2Nq6dl2lwXrIlkPsBve/AckBcalmDwh -tJgAOHM8VEMp0imJ8+SPGZupiIt+sTVSt4aq2uB5dtTBDcXGx4gVUYJoekI4yQtv -OWKAQAobq0Cutogx8hyXr2tSizp5pGpJDKx8/G/3YcNmRqCq+HQgwshlHQy/hPqs -QC/jcuOkr5GOZASF3wLCuLmAwFCELL96iVEJvPLe5Yb1u4HbCrsy9/N8rHV1s1qa -xMcGiIS4m/fwLbD+TjgM1foESeQciHH7GWVZ3Osn6Inpp6vLJKDhLhGBgpaYF72S -0fft6CkIy/CqbwXPe/3Wm0PL0SLTsSk0lKhw6Zeuvg0Z0jUrmDmlTrGo8j4zU7cA -Gc8uiZuFAU2ZoS9R7SZXAoIBAQDSfF8H4C8jn9IvTw/S3APkP7VyLKFSTitEfagt -qcAZr7sYWYGdYDL1hX+jBopAWdzbJ+N3MZBccH67f3XDZWVH4GDQWF6tyOpyseBe -42aU/yCr6ZSoxi91W9V+oyiTFh+t2cEkNyGey/pbXVBw45jncHpzIAaERm5gD8DD -MxOadIvDxEGquGs9MUgZg6ABFmqvQIfiZ7I3PvfoHJGrHsPyM9i6AqAkJe1rKiN3 -tt2wiDk6zDor4vFKn+5LaRQ5INbpgBuCJrvDsRh3MeKLb0ddCYcAStHeAjHJZbH4 -OiC5y9hpaIcSQN3Z4gEhdMd+Lhi1+z36uN5k7KCvbGmO3sG/AoIBAQDKvMiVqTgm -a36FFswwZJd5o+Hx6G8Nt5y7ZYLJgoq4wyTPYA5ZIHBp2Jtb//IGJMKpeLcRigxx -nKu4zrLZPHesgYzHbCH9i5fOBFl8yF9pZnbCU/zGOYhiVdn7MHZgJw8nlstWr1Jr -cMnxBBXEjgL65Lrse1noiPATNrvBsirFaJky6HhxRPbkijKKd0Xqc0FeHMxt6IHU -y6yZRzguI0M1A8RGR4CrqsOGrdv0vkMSiumkUdz8JW4w9R69n6ax/qNAZSMql0ss -PIPOqGtYRPhibOhXKPl0p30X32YXx/SKm8+L9Sr1ny76dab/bSnxStOdihGJLCs4 -l7vFkuJgTMtbAoIBAQCYA3CSfIsu3EbtGdlgvLsmxggh7C+aBJBlB6dFSzpMksi5 -njLo2MgU35Q9xgRk00GZGWbC94295RTyDuya8IjD7z2cWqYONnNz4BkeDndQli0f -WzOc7Hzr8iXvLqCoEatRYFmH8TUbvU8TWwI0dXtBcs9Mg82RDFi8kcPyddnri85A -1WVjiYsRh5z9qD0PbAQii6VXkvJ3ycc64B8oCbEUI/Oa6ziCws2DvswcsnnK+6bx -WvuMJHuFHJn55mrPk3MC8h1r0tN6UlVMCEAH2ZcdjzrrsB1/i/Au9n4gusJVzO1/ -uxkJysUujXWplvBYpav9CfVKNOeQ1gB6kP5vS1t7AoIBAB1DquCPkJ9bHOQxKkBC -BOt2EINOvdkJDAKw4HQd99A7uvCEOQ38dL2Smrpo85KXc9HqruJFPw6XQuJmU8Kv -y8aG3L9ciHuEzuDaF+C/O6aHN9VNMkuaumkXY2Oy1yOB/9oDFk7o98iyezPjFxFM -Pnng0mqYU54RRjY/zFJlWW8tbg+/JsOS5OCQYkNCfEEfaewf1BJ5YWRKEhv9/8oJ -JQZeCNLsN1KQT7D9H6bwX9YpXxhtCK0M6h7/AvT0OqeuzfnZn33iYON9yLjn7rbL -Hd93QQJz065XDuOHR8FfB5mKbCcTuKPD2pAks3pjU46U8n7nEyjtyz9cB6q5TRwB -eckCggEBAMd/3riFnoc2VF6LsA0VQzZiJleN0G9g3AF/5acqHHjMv6aJs/AjV7ZF -hFdiqTcQFEg87HSOzrtV2DVNWJc7DMcMIuwGhVec+D4wedJ2wrQv+MEgySltSBZq -wcPVn5IQiml38NnG/tPIrHETb0PIoa8+iu/Jg80o7j8M3+DKVKfCyfh334PjFK6W -B/mkgC9PfcfeA/Doby9pJsRnqmAJTeWjxbefksckI4PcRCLMEwggB2ReiIWef8Q+ -IooNIxypWtBWtpNEl5lhO7Y2f65Whp34TjooXmGBl1lj4szO8PNnf9QA865J86OS -kOoFda4Sn7LajkbSX0wTGMuXDpmx34M= ------END PRIVATE KEY----- \ No newline at end of file diff --git a/presto-base-jdbc/src/main/java/com/facebook/presto/plugin/jdbc/JdbcMetadata.java b/presto-base-jdbc/src/main/java/com/facebook/presto/plugin/jdbc/JdbcMetadata.java index b56e81b2912a1..9c37bd8a3d731 100644 --- a/presto-base-jdbc/src/main/java/com/facebook/presto/plugin/jdbc/JdbcMetadata.java +++ b/presto-base-jdbc/src/main/java/com/facebook/presto/plugin/jdbc/JdbcMetadata.java @@ -85,11 +85,15 @@ public JdbcTableHandle getTableHandle(ConnectorSession session, SchemaTableName } @Override - public List getTableLayouts(ConnectorSession session, ConnectorTableHandle table, Constraint constraint, Optional> desiredColumns) + public ConnectorTableLayoutResult getTableLayoutForConstraint( + ConnectorSession session, + ConnectorTableHandle table, + Constraint constraint, + Optional> desiredColumns) { JdbcTableHandle tableHandle = (JdbcTableHandle) table; ConnectorTableLayout layout = new ConnectorTableLayout(new JdbcTableLayoutHandle(session.getSqlFunctionProperties(), tableHandle, constraint.getSummary(), Optional.empty())); - return ImmutableList.of(new ConnectorTableLayoutResult(layout, constraint.getSummary())); + return new ConnectorTableLayoutResult(layout, constraint.getSummary()); } @Override diff --git a/presto-base-jdbc/src/main/java/com/facebook/presto/plugin/jdbc/mapping/StandardColumnMappings.java b/presto-base-jdbc/src/main/java/com/facebook/presto/plugin/jdbc/mapping/StandardColumnMappings.java index 3d705e4a2e94e..ed4d18eab0cfa 100644 --- a/presto-base-jdbc/src/main/java/com/facebook/presto/plugin/jdbc/mapping/StandardColumnMappings.java +++ b/presto-base-jdbc/src/main/java/com/facebook/presto/plugin/jdbc/mapping/StandardColumnMappings.java @@ -396,6 +396,15 @@ else if (type instanceof DateType) { else if (type instanceof TimestampType) { return Optional.of(timestampWriteMapping((TimestampType) type)); } + else if (type.equals(TIME)) { + return Optional.of(timeWriteMapping()); + } + else if (type.equals(TIME_WITH_TIME_ZONE)) { + return Optional.of(timeWithTimeZoneWriteMapping()); + } + else if (type.equals(TIMESTAMP_WITH_TIME_ZONE)) { + return Optional.of(timestampWithTimeZoneWriteMapping()); + } else if (type.equals(UUID)) { return Optional.of(uuidWriteMapping()); } @@ -428,6 +437,9 @@ else if (type.equals(DOUBLE)) { else if (type instanceof CharType || type instanceof VarcharType) { return Optional.of(charWriteMapping()); } + else if (type instanceof DecimalType) { + return Optional.of(decimalWriteMapping((DecimalType) type)); + } else if (type.equals(DateType.DATE)) { return Optional.of(dateWriteMapping()); } diff --git a/presto-base-jdbc/src/test/java/com/facebook/presto/plugin/jdbc/TestJdbcDistributedQueries.java b/presto-base-jdbc/src/test/java/com/facebook/presto/plugin/jdbc/TestJdbcDistributedQueries.java index 2d2962742a6e3..8b73ea66dc137 100644 --- a/presto-base-jdbc/src/test/java/com/facebook/presto/plugin/jdbc/TestJdbcDistributedQueries.java +++ b/presto-base-jdbc/src/test/java/com/facebook/presto/plugin/jdbc/TestJdbcDistributedQueries.java @@ -13,11 +13,14 @@ */ package com.facebook.presto.plugin.jdbc; +import com.facebook.presto.Session; import com.facebook.presto.testing.QueryRunner; import com.facebook.presto.tests.AbstractTestQueries; import io.airlift.tpch.TpchTable; +import org.testng.annotations.Test; import static com.facebook.presto.plugin.jdbc.JdbcQueryRunner.createJdbcQueryRunner; +import static com.facebook.presto.testing.TestingSession.testSessionBuilder; public class TestJdbcDistributedQueries extends AbstractTestQueries @@ -33,4 +36,15 @@ protected QueryRunner createQueryRunner() public void testLargeIn() { } + + @Test + public void testNativeQueryParameters() + { + Session session = testSessionBuilder() + .addPreparedStatement("my_query_simple", "SELECT * FROM TABLE(system.query(query => ?))") + .addPreparedStatement("my_query", "SELECT * FROM TABLE(system.query(query => format('SELECT %s FROM %s', ?, ?)))") + .build(); + assertQueryFails(session, "EXECUTE my_query_simple USING 'SELECT 1 a'", "line 1:21: Table function system.query not registered"); + assertQueryFails(session, "EXECUTE my_query USING 'a', '(SELECT 2 a) t'", "line 1:21: Table function system.query not registered"); + } } diff --git a/presto-benchmark-driver/pom.xml b/presto-benchmark-driver/pom.xml index 1ea86335581be..76725407eb1d5 100644 --- a/presto-benchmark-driver/pom.xml +++ b/presto-benchmark-driver/pom.xml @@ -93,48 +93,58 @@ - - - - org.apache.maven.plugins - maven-shade-plugin - - - package - - shade - + + + executable-jar + + + !skipExecutableJar + + + + + + org.apache.maven.plugins + maven-shade-plugin + + + package + + shade + + + true + executable + + + + ${main-class} + + + + + + + + + + org.skife.maven + really-executable-jar-maven-plugin - true - executable - - - - ${main-class} - - - + -Xmx1G + executable - - - - - - org.skife.maven - really-executable-jar-maven-plugin - - -Xmx1G - executable - - - - package - - really-executable-jar - - - - - - + + + package + + really-executable-jar + + + + + + + + diff --git a/presto-benchto-benchmarks/src/main/java/com/facebook/presto/benchto/benchmarks/Dummy.java b/presto-benchto-benchmarks/src/main/java/com/facebook/presto/benchto/benchmarks/Dummy.java new file mode 100644 index 0000000000000..e5bc3b155b6b8 --- /dev/null +++ b/presto-benchto-benchmarks/src/main/java/com/facebook/presto/benchto/benchmarks/Dummy.java @@ -0,0 +1,21 @@ +/* + * Licensed 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 com.facebook.presto.benchto.benchmarks; + +/** + * This class exists to force the creation of a jar for the presto-benchto-benchmarks module. This is needed to deploy the presto-benchto-benchmarks module to nexus. + */ +public class Dummy +{ +} diff --git a/presto-bigquery/pom.xml b/presto-bigquery/pom.xml index 07ad6c02c7b80..611d8703d3226 100644 --- a/presto-bigquery/pom.xml +++ b/presto-bigquery/pom.xml @@ -342,7 +342,7 @@ org.objenesis objenesis - 2.6 + 3.4 test diff --git a/presto-bigquery/src/main/java/com/facebook/presto/plugin/bigquery/BigQueryClient.java b/presto-bigquery/src/main/java/com/facebook/presto/plugin/bigquery/BigQueryClient.java index 59215dc81fe79..221ece8ae3b74 100644 --- a/presto-bigquery/src/main/java/com/facebook/presto/plugin/bigquery/BigQueryClient.java +++ b/presto-bigquery/src/main/java/com/facebook/presto/plugin/bigquery/BigQueryClient.java @@ -50,6 +50,7 @@ public class BigQueryClient private final Optional viewMaterializationProject; private final Optional viewMaterializationDataset; private final String tablePrefix = "_pbc_"; + private final boolean caseSensitiveNameMatching; // presto converts the dataset and table names to lower case, while BigQuery is case sensitive private final ConcurrentMap tableIds = new ConcurrentHashMap<>(); @@ -60,6 +61,7 @@ public class BigQueryClient this.bigQuery = requireNonNull(bigQuery, "bigQuery is null"); this.viewMaterializationProject = requireNonNull(config.getViewMaterializationProject(), "viewMaterializationProject is null"); this.viewMaterializationDataset = requireNonNull(config.getViewMaterializationDataset(), "viewMaterializationDataset is null"); + this.caseSensitiveNameMatching = config.isCaseSensitiveNameMatching(); } public TableInfo getTable(TableId tableId) @@ -108,7 +110,7 @@ private void addTableMappingIfNeeded(DatasetId datasetID, Table table) private Dataset addDataSetMappingIfNeeded(Dataset dataset) { DatasetId bigQueryDatasetId = dataset.getDatasetId(); - DatasetId prestoDatasetId = DatasetId.of(bigQueryDatasetId.getProject(), bigQueryDatasetId.getDataset().toLowerCase(ENGLISH)); + DatasetId prestoDatasetId = DatasetId.of(bigQueryDatasetId.getProject(), bigQueryDatasetId.getDataset()); datasetIds.putIfAbsent(prestoDatasetId, bigQueryDatasetId); return dataset; } @@ -123,7 +125,8 @@ protected TableId createDestinationTable(TableId tableId) private String createTableName() { - return format(tablePrefix + "%s", randomUUID().toString().toLowerCase(ENGLISH).replace("-", "")); + String uuid = randomUUID().toString().replace("-", ""); + return caseSensitiveNameMatching ? format("%s%s", tablePrefix, uuid) : format("%s%s", tablePrefix, uuid).toLowerCase(ENGLISH); } private DatasetId mapIfNeeded(String project, String dataset) diff --git a/presto-bigquery/src/main/java/com/facebook/presto/plugin/bigquery/BigQueryConfig.java b/presto-bigquery/src/main/java/com/facebook/presto/plugin/bigquery/BigQueryConfig.java index 86e09dba9077e..1247f385e1514 100644 --- a/presto-bigquery/src/main/java/com/facebook/presto/plugin/bigquery/BigQueryConfig.java +++ b/presto-bigquery/src/main/java/com/facebook/presto/plugin/bigquery/BigQueryConfig.java @@ -35,6 +35,7 @@ public class BigQueryConfig private Optional parentProjectId = Optional.empty(); private OptionalInt parallelism = OptionalInt.empty(); private boolean viewsEnabled; + private boolean caseSensitiveNameMatching; private Optional viewMaterializationProject = Optional.empty(); private Optional viewMaterializationDataset = Optional.empty(); private int maxReadRowsRetries = DEFAULT_MAX_READ_ROWS_RETRIES; @@ -181,6 +182,22 @@ public BigQueryConfig setMaxReadRowsRetries(int maxReadRowsRetries) return this; } + public boolean isCaseSensitiveNameMatching() + { + return caseSensitiveNameMatching; + } + + @Config("case-sensitive-name-matching") + @ConfigDescription( + "Case sensitivity for schema and table name matching. " + + "true = preserve case and require exact matches; " + + "false (default) = normalize to lower case and match case-insensitively.") + public BigQueryConfig setCaseSensitiveNameMatching(boolean caseSensitiveNameMatching) + { + this.caseSensitiveNameMatching = caseSensitiveNameMatching; + return this; + } + ReadSessionCreatorConfig createReadSessionCreatorConfig() { return new ReadSessionCreatorConfig( diff --git a/presto-bigquery/src/main/java/com/facebook/presto/plugin/bigquery/BigQueryMetadata.java b/presto-bigquery/src/main/java/com/facebook/presto/plugin/bigquery/BigQueryMetadata.java index 8c028fbf59b7f..45dfb474efb81 100644 --- a/presto-bigquery/src/main/java/com/facebook/presto/plugin/bigquery/BigQueryMetadata.java +++ b/presto-bigquery/src/main/java/com/facebook/presto/plugin/bigquery/BigQueryMetadata.java @@ -51,6 +51,7 @@ import static com.google.cloud.bigquery.TableDefinition.Type.TABLE; import static com.google.cloud.bigquery.TableDefinition.Type.VIEW; import static com.google.common.collect.ImmutableList.toImmutableList; +import static java.util.Locale.ENGLISH; import static java.util.Objects.requireNonNull; import static java.util.stream.Collectors.toMap; @@ -63,12 +64,14 @@ public class BigQueryMetadata private static final Logger log = Logger.get(BigQueryMetadata.class); private final BigQueryClient bigQueryClient; private final String projectId; + private final boolean caseSensitiveNameMatching; @Inject public BigQueryMetadata(BigQueryClient bigQueryClient, BigQueryConfig config) { this.bigQueryClient = bigQueryClient; this.projectId = config.getProjectId().orElse(bigQueryClient.getProjectId()); + this.caseSensitiveNameMatching = config.isCaseSensitiveNameMatching(); } @Override @@ -124,7 +127,7 @@ public ConnectorTableHandle getTableHandle(ConnectorSession session, SchemaTable } @Override - public List getTableLayouts( + public ConnectorTableLayoutResult getTableLayoutForConstraint( ConnectorSession session, ConnectorTableHandle table, Constraint constraint, @@ -136,7 +139,7 @@ public List getTableLayouts( bigQueryTableHandle = bigQueryTableHandle.withProjectedColumns(ImmutableList.copyOf(desiredColumns.get())); } BigQueryTableLayoutHandle bigQueryTableLayoutHandle = new BigQueryTableLayoutHandle(bigQueryTableHandle); - return ImmutableList.of(new ConnectorTableLayoutResult(new ConnectorTableLayout(bigQueryTableLayoutHandle), constraint.getSummary())); + return new ConnectorTableLayoutResult(new ConnectorTableLayout(bigQueryTableLayoutHandle), constraint.getSummary()); } @Override @@ -233,4 +236,10 @@ private List listTables(ConnectorSession session, SchemaTablePr ImmutableList.of(tableName) : ImmutableList.of(); // table does not exist } + + @Override + public String normalizeIdentifier(ConnectorSession session, String identifier) + { + return caseSensitiveNameMatching ? identifier : identifier.toLowerCase(ENGLISH); + } } diff --git a/presto-bigquery/src/test/java/com/facebook/presto/plugin/bigquery/TestBigQueryConfig.java b/presto-bigquery/src/test/java/com/facebook/presto/plugin/bigquery/TestBigQueryConfig.java index 053f83600eb85..8829d86cd6e23 100644 --- a/presto-bigquery/src/test/java/com/facebook/presto/plugin/bigquery/TestBigQueryConfig.java +++ b/presto-bigquery/src/test/java/com/facebook/presto/plugin/bigquery/TestBigQueryConfig.java @@ -36,7 +36,8 @@ public void testDefaults() .setParallelism(20) .setViewMaterializationProject("vmproject") .setViewMaterializationDataset("vmdataset") - .setMaxReadRowsRetries(10); + .setMaxReadRowsRetries(10) + .setCaseSensitiveNameMatching(false); assertEquals(config.getCredentialsKey(), Optional.of("ckey")); assertEquals(config.getCredentialsFile(), Optional.of("cfile")); @@ -46,6 +47,7 @@ public void testDefaults() assertEquals(config.getViewMaterializationProject(), Optional.of("vmproject")); assertEquals(config.getViewMaterializationDataset(), Optional.of("vmdataset")); assertEquals(config.getMaxReadRowsRetries(), 10); + assertEquals(config.isCaseSensitiveNameMatching(), false); } @Test @@ -59,6 +61,7 @@ public void testExplicitPropertyMappingsWithCredentialsKey() .put("bigquery.view-materialization-project", "vmproject") .put("bigquery.view-materialization-dataset", "vmdataset") .put("bigquery.max-read-rows-retries", "10") + .put("case-sensitive-name-matching", "true") .build(); ConfigurationFactory configurationFactory = new ConfigurationFactory(properties); @@ -71,6 +74,7 @@ public void testExplicitPropertyMappingsWithCredentialsKey() assertEquals(config.getViewMaterializationProject(), Optional.of("vmproject")); assertEquals(config.getViewMaterializationDataset(), Optional.of("vmdataset")); assertEquals(config.getMaxReadRowsRetries(), 10); + assertEquals(config.isCaseSensitiveNameMatching(), true); } @Test diff --git a/presto-blackhole/src/main/java/com/facebook/presto/plugin/blackhole/BlackHoleMetadata.java b/presto-blackhole/src/main/java/com/facebook/presto/plugin/blackhole/BlackHoleMetadata.java index 97ea968c5bbdc..afe0f1efd8dc4 100644 --- a/presto-blackhole/src/main/java/com/facebook/presto/plugin/blackhole/BlackHoleMetadata.java +++ b/presto-blackhole/src/main/java/com/facebook/presto/plugin/blackhole/BlackHoleMetadata.java @@ -253,7 +253,7 @@ public Optional finishInsert(ConnectorSession session, } @Override - public List getTableLayouts( + public ConnectorTableLayoutResult getTableLayoutForConstraint( ConnectorSession session, ConnectorTableHandle handle, Constraint constraint, @@ -266,7 +266,7 @@ public List getTableLayouts( blackHoleHandle.getRowsPerPage(), blackHoleHandle.getFieldsLength(), blackHoleHandle.getPageProcessingDelay()); - return ImmutableList.of(new ConnectorTableLayoutResult(getTableLayout(session, layoutHandle), constraint.getSummary())); + return new ConnectorTableLayoutResult(getTableLayout(session, layoutHandle), constraint.getSummary()); } @Override diff --git a/presto-built-in-worker-function-tools/pom.xml b/presto-built-in-worker-function-tools/pom.xml new file mode 100644 index 0000000000000..8159c62c45b5d --- /dev/null +++ b/presto-built-in-worker-function-tools/pom.xml @@ -0,0 +1,62 @@ + + + + presto-root + com.facebook.presto + 0.295-SNAPSHOT + + 4.0.0 + + presto-built-in-worker-function-tools + presto-built-in-worker-function-tools + + + ${project.parent.basedir} + 17 + true + + + + + com.facebook.presto + presto-spi + + + com.facebook.presto + presto-function-namespace-managers-common + + + com.facebook.airlift + http-client + + + com.google.inject + guice + + + com.google.guava + guava + + + com.facebook.airlift + json + + + com.facebook.presto + presto-common + + + com.facebook.airlift + log + + + com.facebook.airlift + configuration + + + org.testng + testng + test + + + diff --git a/presto-built-in-worker-function-tools/src/main/java/com/facebook/presto/builtin/tools/ForNativeFunctionRegistryInfo.java b/presto-built-in-worker-function-tools/src/main/java/com/facebook/presto/builtin/tools/ForNativeFunctionRegistryInfo.java new file mode 100644 index 0000000000000..000a6b2e6d0e1 --- /dev/null +++ b/presto-built-in-worker-function-tools/src/main/java/com/facebook/presto/builtin/tools/ForNativeFunctionRegistryInfo.java @@ -0,0 +1,26 @@ +/* + * Licensed 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 com.facebook.presto.builtin.tools; + +import com.google.inject.BindingAnnotation; + +import java.lang.annotation.Retention; + +import static java.lang.annotation.RetentionPolicy.RUNTIME; + +@Retention(RUNTIME) +@BindingAnnotation +public @interface ForNativeFunctionRegistryInfo +{ +} diff --git a/presto-built-in-worker-function-tools/src/main/java/com/facebook/presto/builtin/tools/NativeSidecarFunctionRegistryTool.java b/presto-built-in-worker-function-tools/src/main/java/com/facebook/presto/builtin/tools/NativeSidecarFunctionRegistryTool.java new file mode 100644 index 0000000000000..1990918df633f --- /dev/null +++ b/presto-built-in-worker-function-tools/src/main/java/com/facebook/presto/builtin/tools/NativeSidecarFunctionRegistryTool.java @@ -0,0 +1,120 @@ +/* + * Licensed 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 com.facebook.presto.builtin.tools; + +import com.facebook.airlift.http.client.HttpClient; +import com.facebook.airlift.http.client.HttpUriBuilder; +import com.facebook.airlift.http.client.JsonResponseHandler; +import com.facebook.airlift.http.client.Request; +import com.facebook.airlift.json.JsonCodec; +import com.facebook.airlift.log.Logger; +import com.facebook.presto.functionNamespace.JsonBasedUdfFunctionMetadata; +import com.facebook.presto.functionNamespace.UdfFunctionSignatureMap; +import com.facebook.presto.spi.Node; +import com.facebook.presto.spi.NodeManager; +import com.facebook.presto.spi.PrestoException; +import com.facebook.presto.spi.StandardErrorCode; +import com.facebook.presto.spi.function.SqlFunction; +import com.google.common.collect.ImmutableMap; + +import java.net.URI; +import java.util.List; +import java.util.Map; + +import static com.google.common.collect.ImmutableList.toImmutableList; +import static java.util.Objects.requireNonNull; + +public class NativeSidecarFunctionRegistryTool + implements WorkerFunctionRegistryTool +{ + private final int maxRetries; + private final long retryDelayMs; + private static final Logger log = Logger.get(NativeSidecarFunctionRegistryTool.class); + private final JsonCodec>> nativeFunctionSignatureMapJsonCodec; + private final NodeManager nodeManager; + private final HttpClient httpClient; + private static final String FUNCTION_SIGNATURES_ENDPOINT = "/v1/functions"; + + public NativeSidecarFunctionRegistryTool( + HttpClient httpClient, + JsonCodec>> nativeFunctionSignatureMapJsonCodec, + NodeManager nodeManager, + int nativeSidecarRegistryToolNumRetries, + long nativeSidecarRegistryToolRetryDelayMs) + { + this.nativeFunctionSignatureMapJsonCodec = + requireNonNull(nativeFunctionSignatureMapJsonCodec, "nativeFunctionSignatureMapJsonCodec is null"); + this.nodeManager = requireNonNull(nodeManager, "nodeManager is null"); + this.httpClient = requireNonNull(httpClient, "typeManager is null"); + this.maxRetries = nativeSidecarRegistryToolNumRetries; + this.retryDelayMs = nativeSidecarRegistryToolRetryDelayMs; + } + + @Override + public List getWorkerFunctions() + { + return getNativeFunctionSignatureMap() + .getUDFSignatureMap() + .entrySet() + .stream() + .flatMap(entry -> entry.getValue().stream() + .map(metaInfo -> WorkerFunctionUtil.createSqlInvokedFunction(entry.getKey(), metaInfo, "presto"))) + .collect(toImmutableList()); + } + + private UdfFunctionSignatureMap getNativeFunctionSignatureMap() + { + try { + Request request = Request.Builder.prepareGet().setUri(getSidecarLocationOnStartup(nodeManager, maxRetries, retryDelayMs)).build(); + Map> nativeFunctionSignatureMap = httpClient.execute(request, JsonResponseHandler.createJsonResponseHandler(nativeFunctionSignatureMapJsonCodec)); + return new UdfFunctionSignatureMap(ImmutableMap.copyOf(nativeFunctionSignatureMap)); + } + catch (Exception e) { + throw new PrestoException(StandardErrorCode.INVALID_ARGUMENTS, "Failed to get functions from sidecar.", e); + } + } + + public static URI getSidecarLocationOnStartup(NodeManager nodeManager, int maxRetries, long retryDelayMs) + { + Node sidecarNode = null; + for (int attempt = 1; attempt <= maxRetries; attempt++) { + try { + sidecarNode = nodeManager.getSidecarNode(); + if (sidecarNode != null) { + break; + } + } + catch (Exception e) { + log.error("Error getting sidecar node (attempt " + attempt + "): " + e.getMessage()); + if (attempt == maxRetries) { + throw new RuntimeException("Failed to get sidecar node", e); + } + else { + try { + Thread.sleep(retryDelayMs); + } + catch (InterruptedException ie) { + Thread.currentThread().interrupt(); + throw new RuntimeException("Retry fetching sidecar function registry interrupted", ie); + } + } + } + } + + return HttpUriBuilder + .uriBuilderFrom(sidecarNode.getHttpUri()) + .appendPath(FUNCTION_SIGNATURES_ENDPOINT) + .build(); + } +} diff --git a/presto-built-in-worker-function-tools/src/main/java/com/facebook/presto/builtin/tools/NativeSidecarRegistryToolConfig.java b/presto-built-in-worker-function-tools/src/main/java/com/facebook/presto/builtin/tools/NativeSidecarRegistryToolConfig.java new file mode 100644 index 0000000000000..4ce944572406d --- /dev/null +++ b/presto-built-in-worker-function-tools/src/main/java/com/facebook/presto/builtin/tools/NativeSidecarRegistryToolConfig.java @@ -0,0 +1,51 @@ +/* + * Licensed 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 com.facebook.presto.builtin.tools; + +import com.facebook.airlift.configuration.Config; +import com.facebook.airlift.configuration.ConfigDescription; + +import java.time.Duration; + +public class NativeSidecarRegistryToolConfig +{ + private int nativeSidecarRegistryToolNumRetries = 8; + private long nativeSidecarRegistryToolRetryDelayMs = Duration.ofMinutes(1).toMillis(); + + public int getNativeSidecarRegistryToolNumRetries() + { + return nativeSidecarRegistryToolNumRetries; + } + + @Config("native-sidecar-registry-tool.num-retries") + @ConfigDescription("Max times to retry fetching sidecar node") + public NativeSidecarRegistryToolConfig setNativeSidecarRegistryToolNumRetries(int nativeSidecarRegistryToolNumRetries) + { + this.nativeSidecarRegistryToolNumRetries = nativeSidecarRegistryToolNumRetries; + return this; + } + + public long getNativeSidecarRegistryToolRetryDelayMs() + { + return nativeSidecarRegistryToolRetryDelayMs; + } + + @Config("native-sidecar-registry-tool.retry-delay-ms") + @ConfigDescription("Cooldown period to retry when fetching sidecar node fails") + public NativeSidecarRegistryToolConfig setNativeSidecarRegistryToolRetryDelayMs(long nativeSidecarRegistryToolRetryDelayMs) + { + this.nativeSidecarRegistryToolRetryDelayMs = nativeSidecarRegistryToolRetryDelayMs; + return this; + } +} diff --git a/presto-built-in-worker-function-tools/src/main/java/com/facebook/presto/builtin/tools/WorkerFunctionRegistryTool.java b/presto-built-in-worker-function-tools/src/main/java/com/facebook/presto/builtin/tools/WorkerFunctionRegistryTool.java new file mode 100644 index 0000000000000..99c7cfe13bc78 --- /dev/null +++ b/presto-built-in-worker-function-tools/src/main/java/com/facebook/presto/builtin/tools/WorkerFunctionRegistryTool.java @@ -0,0 +1,23 @@ +/* + * Licensed 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 com.facebook.presto.builtin.tools; + +import com.facebook.presto.spi.function.SqlFunction; + +import java.util.List; + +public interface WorkerFunctionRegistryTool +{ + List getWorkerFunctions(); +} diff --git a/presto-built-in-worker-function-tools/src/main/java/com/facebook/presto/builtin/tools/WorkerFunctionUtil.java b/presto-built-in-worker-function-tools/src/main/java/com/facebook/presto/builtin/tools/WorkerFunctionUtil.java new file mode 100644 index 0000000000000..d880429f7489f --- /dev/null +++ b/presto-built-in-worker-function-tools/src/main/java/com/facebook/presto/builtin/tools/WorkerFunctionUtil.java @@ -0,0 +1,175 @@ +/* + * Licensed 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 com.facebook.presto.builtin.tools; + +import com.facebook.presto.common.CatalogSchemaName; +import com.facebook.presto.common.QualifiedObjectName; +import com.facebook.presto.common.type.NamedTypeSignature; +import com.facebook.presto.common.type.StandardTypes; +import com.facebook.presto.common.type.TypeSignature; +import com.facebook.presto.common.type.TypeSignatureParameter; +import com.facebook.presto.functionNamespace.JsonBasedUdfFunctionMetadata; +import com.facebook.presto.spi.function.AggregationFunctionMetadata; +import com.facebook.presto.spi.function.LongVariableConstraint; +import com.facebook.presto.spi.function.Parameter; +import com.facebook.presto.spi.function.RoutineCharacteristics; +import com.facebook.presto.spi.function.SqlInvokedFunction; +import com.facebook.presto.spi.function.TypeVariableConstraint; +import com.google.common.collect.ImmutableList; + +import java.util.ArrayList; +import java.util.List; +import java.util.Optional; + +import static com.facebook.presto.spi.function.FunctionVersion.notVersioned; +import static com.google.common.base.Preconditions.checkArgument; +import static com.google.common.base.Preconditions.checkState; + +public class WorkerFunctionUtil +{ + private WorkerFunctionUtil() {} + + public static synchronized SqlInvokedFunction createSqlInvokedFunction(String functionName, JsonBasedUdfFunctionMetadata jsonBasedUdfFunctionMetaData, String catalogName) + { + checkState(jsonBasedUdfFunctionMetaData.getRoutineCharacteristics().getLanguage().equals(RoutineCharacteristics.Language.CPP), "WorkerFunctionUtil only supports CPP UDF"); + QualifiedObjectName qualifiedFunctionName = QualifiedObjectName.valueOf(new CatalogSchemaName(catalogName, jsonBasedUdfFunctionMetaData.getSchema()), functionName); + List parameterNameList = jsonBasedUdfFunctionMetaData.getParamNames(); + List parameterTypeList = convertApplicableTypeToVariable(jsonBasedUdfFunctionMetaData.getParamTypes()); + List typeVariableConstraintsList = jsonBasedUdfFunctionMetaData.getTypeVariableConstraints().isPresent() ? + jsonBasedUdfFunctionMetaData.getTypeVariableConstraints().get() : ImmutableList.of(); + List longVariableConstraintList = jsonBasedUdfFunctionMetaData.getLongVariableConstraints().isPresent() ? + jsonBasedUdfFunctionMetaData.getLongVariableConstraints().get() : ImmutableList.of(); + + TypeSignature outputType = convertApplicableTypeToVariable(jsonBasedUdfFunctionMetaData.getOutputType()); + ImmutableList.Builder parameterBuilder = ImmutableList.builder(); + for (int i = 0; i < parameterNameList.size(); i++) { + parameterBuilder.add(new Parameter(parameterNameList.get(i), parameterTypeList.get(i))); + } + + Optional aggregationFunctionMetadata = + jsonBasedUdfFunctionMetaData.getAggregateMetadata() + .map(metadata -> new AggregationFunctionMetadata( + convertApplicableTypeToVariable(metadata.getIntermediateType()), + metadata.isOrderSensitive())); + + return new SqlInvokedFunction( + qualifiedFunctionName, + parameterBuilder.build(), + typeVariableConstraintsList, + longVariableConstraintList, + outputType, + jsonBasedUdfFunctionMetaData.getDocString(), + jsonBasedUdfFunctionMetaData.getRoutineCharacteristics(), + "", + jsonBasedUdfFunctionMetaData.getVariableArity(), + notVersioned(), + jsonBasedUdfFunctionMetaData.getFunctionKind(), + aggregationFunctionMetadata); + } + + // Todo: Improve the handling of parameter type differentiation in native execution. + // HACK: Currently, we lack support for correctly identifying the parameterKind, specifically between TYPE and VARIABLE, + // in native execution. The following utility functions help bridge this gap by parsing the type signature and verifying whether its base + // and parameters are of a supported type. The valid types list are non - parametric types that Presto supports. + public static List convertApplicableTypeToVariable(List typeSignatures) + { + List newTypeSignaturesList = new ArrayList<>(); + for (TypeSignature typeSignature : typeSignatures) { + if (!typeSignature.getParameters().isEmpty()) { + TypeSignature newTypeSignature = + new TypeSignature( + typeSignature.getBase(), + getTypeSignatureParameters( + typeSignature, + typeSignature.getParameters())); + newTypeSignaturesList.add(newTypeSignature); + } + else { + newTypeSignaturesList.add(typeSignature); + } + } + return newTypeSignaturesList; + } + + public static TypeSignature convertApplicableTypeToVariable(TypeSignature typeSignature) + { + List typeSignaturesList = convertApplicableTypeToVariable(ImmutableList.of(typeSignature)); + checkArgument(!typeSignaturesList.isEmpty(), "Type signature list is empty for : " + typeSignature); + return typeSignaturesList.get(0); + } + + private static List getTypeSignatureParameters( + TypeSignature typeSignature, + List typeSignatureParameterList) + { + List newParameterTypeList = new ArrayList<>(); + for (TypeSignatureParameter parameter : typeSignatureParameterList) { + if (parameter.isLongLiteral()) { + newParameterTypeList.add(parameter); + continue; + } + + boolean isNamedTypeSignature = parameter.isNamedTypeSignature(); + TypeSignature parameterTypeSignature; + // If it's a named type signatures only in the case of row signature types. + if (isNamedTypeSignature) { + parameterTypeSignature = parameter.getNamedTypeSignature().getTypeSignature(); + } + else { + parameterTypeSignature = parameter.getTypeSignature(); + } + + if (parameterTypeSignature.getParameters().isEmpty()) { + boolean changeTypeToVariable = isDecimalTypeBase(typeSignature.getBase()); + if (changeTypeToVariable) { + newParameterTypeList.add( + TypeSignatureParameter.of(parameterTypeSignature.getBase())); + } + else { + if (isNamedTypeSignature) { + newParameterTypeList.add(TypeSignatureParameter.of(parameter.getNamedTypeSignature())); + } + else { + newParameterTypeList.add(TypeSignatureParameter.of(parameterTypeSignature)); + } + } + } + else { + TypeSignature newTypeSignature = + new TypeSignature( + parameterTypeSignature.getBase(), + getTypeSignatureParameters( + parameterTypeSignature.getStandardTypeSignature(), + parameterTypeSignature.getParameters())); + if (isNamedTypeSignature) { + newParameterTypeList.add( + TypeSignatureParameter.of( + new NamedTypeSignature( + Optional.empty(), + newTypeSignature))); + } + else { + newParameterTypeList.add(TypeSignatureParameter.of(newTypeSignature)); + } + } + } + return newParameterTypeList; + } + + private static boolean isDecimalTypeBase(String typeBase) + { + return typeBase.equals(StandardTypes.DECIMAL); + } +} diff --git a/presto-built-in-worker-function-tools/src/test/java/com/facebook/presto/builtin/tools/TestNativeSidecarRegistryToolConfig.java b/presto-built-in-worker-function-tools/src/test/java/com/facebook/presto/builtin/tools/TestNativeSidecarRegistryToolConfig.java new file mode 100644 index 0000000000000..34df3f89b88f3 --- /dev/null +++ b/presto-built-in-worker-function-tools/src/test/java/com/facebook/presto/builtin/tools/TestNativeSidecarRegistryToolConfig.java @@ -0,0 +1,50 @@ +/* + * Licensed 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 com.facebook.presto.builtin.tools; + +import com.google.common.collect.ImmutableMap; +import org.testng.annotations.Test; + +import java.util.Map; + +import static com.facebook.airlift.configuration.testing.ConfigAssertions.assertFullMapping; +import static com.facebook.airlift.configuration.testing.ConfigAssertions.assertRecordedDefaults; +import static com.facebook.airlift.configuration.testing.ConfigAssertions.recordDefaults; + +public class TestNativeSidecarRegistryToolConfig +{ + @Test + public void testDefaults() + { + assertRecordedDefaults(recordDefaults(NativeSidecarRegistryToolConfig.class) + .setNativeSidecarRegistryToolNumRetries(8) + .setNativeSidecarRegistryToolRetryDelayMs(60_000L)); + } + + @Test + public void testExplicitPropertyMappings() + throws Exception + { + Map properties = new ImmutableMap.Builder() + .put("native-sidecar-registry-tool.num-retries", "15") + .put("native-sidecar-registry-tool.retry-delay-ms", "11115") + .build(); + + NativeSidecarRegistryToolConfig expected = new NativeSidecarRegistryToolConfig() + .setNativeSidecarRegistryToolNumRetries(15) + .setNativeSidecarRegistryToolRetryDelayMs(11_115L); + + assertFullMapping(properties, expected); + } +} diff --git a/presto-cassandra/src/main/java/com/facebook/presto/cassandra/CassandraMetadata.java b/presto-cassandra/src/main/java/com/facebook/presto/cassandra/CassandraMetadata.java index f609fb22a47ff..e393803401b95 100644 --- a/presto-cassandra/src/main/java/com/facebook/presto/cassandra/CassandraMetadata.java +++ b/presto-cassandra/src/main/java/com/facebook/presto/cassandra/CassandraMetadata.java @@ -201,7 +201,11 @@ public ColumnMetadata getColumnMetadata(ConnectorSession session, ConnectorTable } @Override - public List getTableLayouts(ConnectorSession session, ConnectorTableHandle table, Constraint constraint, Optional> desiredColumns) + public ConnectorTableLayoutResult getTableLayoutForConstraint( + ConnectorSession session, + ConnectorTableHandle table, + Constraint constraint, + Optional> desiredColumns) { CassandraTableHandle handle = (CassandraTableHandle) table; CassandraPartitionResult partitionResult = partitionManager.getPartitions(handle, constraint.getSummary()); @@ -224,7 +228,7 @@ public List getTableLayouts(ConnectorSession session handle, partitionResult.getPartitions(), clusteringKeyPredicates)); - return ImmutableList.of(new ConnectorTableLayoutResult(layout, unenforcedConstraint)); + return new ConnectorTableLayoutResult(layout, unenforcedConstraint); } @Override diff --git a/presto-cassandra/src/test/java/com/facebook/presto/cassandra/TestCassandraConnector.java b/presto-cassandra/src/test/java/com/facebook/presto/cassandra/TestCassandraConnector.java index 268e368918a65..4f5856a09e4e6 100644 --- a/presto-cassandra/src/test/java/com/facebook/presto/cassandra/TestCassandraConnector.java +++ b/presto-cassandra/src/test/java/com/facebook/presto/cassandra/TestCassandraConnector.java @@ -67,7 +67,6 @@ import static com.facebook.presto.spi.connector.ConnectorSplitManager.SplitSchedulingStrategy.UNGROUPED_SCHEDULING; import static com.facebook.presto.spi.connector.NotPartitionedPartitionHandle.NOT_PARTITIONED; import static com.google.common.base.Preconditions.checkArgument; -import static com.google.common.collect.Iterables.getOnlyElement; import static java.util.Locale.ENGLISH; import static org.testng.Assert.assertEquals; import static org.testng.Assert.assertNull; @@ -185,8 +184,8 @@ public void testGetRecords() ConnectorTransactionHandle transaction = CassandraTransactionHandle.INSTANCE; - List layouts = metadata.getTableLayouts(SESSION, tableHandle, Constraint.alwaysTrue(), Optional.empty()); - ConnectorTableLayoutHandle layout = getOnlyElement(layouts).getTableLayout().getHandle(); + ConnectorTableLayoutResult layoutResult = metadata.getTableLayoutForConstraint(SESSION, tableHandle, Constraint.alwaysTrue(), Optional.empty()); + ConnectorTableLayoutHandle layout = layoutResult.getTableLayout().getHandle(); List splits = getAllSplits(splitManager.getSplits(transaction, SESSION, layout, new SplitSchedulingContext(UNGROUPED_SCHEDULING, false, WarningCollector.NOOP))); long rowNumber = 0; diff --git a/presto-cli/pom.xml b/presto-cli/pom.xml index 8d2bcb99858b6..c2298b07486c2 100644 --- a/presto-cli/pom.xml +++ b/presto-cli/pom.xml @@ -130,48 +130,6 @@ - - org.apache.maven.plugins - maven-shade-plugin - - - package - - shade - - - true - executable - false - - - - ${main-class} - - - - - - - - - - org.skife.maven - really-executable-jar-maven-plugin - - -Xmx1G - executable - - - - package - - really-executable-jar - - - - - org.basepom.maven duplicate-finder-maven-plugin @@ -195,4 +153,60 @@ + + + + executable-jar + + + !skipExecutableJar + + + + + + org.apache.maven.plugins + maven-shade-plugin + + + package + + shade + + + true + executable + false + + + + ${main-class} + + + + + + + + + + org.skife.maven + really-executable-jar-maven-plugin + + -Xmx1G + executable + + + + package + + really-executable-jar + + + + + + + + diff --git a/presto-cli/src/main/java/com/facebook/presto/cli/KeyReader.java b/presto-cli/src/main/java/com/facebook/presto/cli/KeyReader.java index 7a75fa24603f3..a242a98511c92 100644 --- a/presto-cli/src/main/java/com/facebook/presto/cli/KeyReader.java +++ b/presto-cli/src/main/java/com/facebook/presto/cli/KeyReader.java @@ -18,11 +18,12 @@ import java.io.IOException; import java.io.InputStream; -import static org.fusesource.jansi.internal.CLibrary.STDIN_FILENO; import static org.fusesource.jansi.internal.CLibrary.isatty; public final class KeyReader { + private static final int STDIN_FILENO = 0; + private KeyReader() {} @SuppressWarnings("resource") diff --git a/presto-clickhouse/src/main/java/com/facebook/presto/plugin/clickhouse/ClickHouseMetadata.java b/presto-clickhouse/src/main/java/com/facebook/presto/plugin/clickhouse/ClickHouseMetadata.java index 6e8fdf91e40ea..749da9652ded4 100755 --- a/presto-clickhouse/src/main/java/com/facebook/presto/plugin/clickhouse/ClickHouseMetadata.java +++ b/presto-clickhouse/src/main/java/com/facebook/presto/plugin/clickhouse/ClickHouseMetadata.java @@ -84,11 +84,15 @@ public ClickHouseTableHandle getTableHandle(ConnectorSession session, SchemaTabl } @Override - public List getTableLayouts(ConnectorSession session, ConnectorTableHandle table, Constraint constraint, Optional> desiredColumns) + public ConnectorTableLayoutResult getTableLayoutForConstraint( + ConnectorSession session, + ConnectorTableHandle table, + Constraint constraint, + Optional> desiredColumns) { ClickHouseTableHandle tableHandle = (ClickHouseTableHandle) table; ConnectorTableLayout layout = new ConnectorTableLayout(new ClickHouseTableLayoutHandle(tableHandle, constraint.getSummary(), Optional.empty(), Optional.empty(), Optional.empty())); - return ImmutableList.of(new ConnectorTableLayoutResult(layout, constraint.getSummary())); + return new ConnectorTableLayoutResult(layout, constraint.getSummary()); } @Override diff --git a/presto-client/pom.xml b/presto-client/pom.xml index e61853a18f6c4..6435a2cda2551 100644 --- a/presto-client/pom.xml +++ b/presto-client/pom.xml @@ -21,7 +21,6 @@ com.facebook.presto presto-spi - @@ -101,19 +100,41 @@ okio-jvm + + com.google.code.findbugs + jsr305 + + org.jetbrains.kotlin kotlin-stdlib-jdk8 + + net.jodah + failsafe + + + + com.facebook.airlift + concurrent + test + + com.facebook.presto presto-testng-services test + + org.assertj + assertj-core + test + + org.testng testng @@ -142,6 +163,12 @@ drift-codec-utils test + + + com.squareup.okhttp3 + mockwebserver + test + diff --git a/presto-client/src/main/java/com/facebook/presto/client/auth/external/CompositeRedirectHandler.java b/presto-client/src/main/java/com/facebook/presto/client/auth/external/CompositeRedirectHandler.java new file mode 100644 index 0000000000000..bc1ebb0355a3a --- /dev/null +++ b/presto-client/src/main/java/com/facebook/presto/client/auth/external/CompositeRedirectHandler.java @@ -0,0 +1,53 @@ +/* + * Licensed 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 com.facebook.presto.client.auth.external; + +import java.net.URI; +import java.util.List; + +import static com.google.common.base.Preconditions.checkState; +import static com.google.common.collect.ImmutableList.toImmutableList; +import static java.util.Objects.requireNonNull; + +public class CompositeRedirectHandler + implements RedirectHandler +{ + private final List handlers; + + public CompositeRedirectHandler(List strategies) + { + this.handlers = requireNonNull(strategies, "strategies is null") + .stream() + .map(ExternalRedirectStrategy::getHandler) + .collect(toImmutableList()); + checkState(!handlers.isEmpty(), "Expected at least one external redirect handler"); + } + + @Override + public void redirectTo(URI uri) throws RedirectException + { + RedirectException redirectException = new RedirectException("Could not redirect to " + uri); + for (RedirectHandler handler : handlers) { + try { + handler.redirectTo(uri); + return; + } + catch (RedirectException e) { + redirectException.addSuppressed(e); + } + } + + throw redirectException; + } +} diff --git a/presto-client/src/main/java/com/facebook/presto/client/auth/external/DesktopBrowserRedirectHandler.java b/presto-client/src/main/java/com/facebook/presto/client/auth/external/DesktopBrowserRedirectHandler.java new file mode 100644 index 0000000000000..6918cba955c4b --- /dev/null +++ b/presto-client/src/main/java/com/facebook/presto/client/auth/external/DesktopBrowserRedirectHandler.java @@ -0,0 +1,41 @@ +/* + * Licensed 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 com.facebook.presto.client.auth.external; + +import java.io.IOException; +import java.net.URI; + +import static java.awt.Desktop.Action.BROWSE; +import static java.awt.Desktop.getDesktop; +import static java.awt.Desktop.isDesktopSupported; + +public final class DesktopBrowserRedirectHandler + implements RedirectHandler +{ + @Override + public void redirectTo(URI uri) + throws RedirectException + { + if (!isDesktopSupported() || !getDesktop().isSupported(BROWSE)) { + throw new RedirectException("Desktop Browser is not available. Make sure your Java process is not in headless mode (-Djava.awt.headless=false)"); + } + + try { + getDesktop().browse(uri); + } + catch (IOException e) { + throw new RedirectException("Failed to redirect", e); + } + } +} diff --git a/presto-client/src/main/java/com/facebook/presto/client/auth/external/ExternalAuthentication.java b/presto-client/src/main/java/com/facebook/presto/client/auth/external/ExternalAuthentication.java new file mode 100644 index 0000000000000..afee9bd899b42 --- /dev/null +++ b/presto-client/src/main/java/com/facebook/presto/client/auth/external/ExternalAuthentication.java @@ -0,0 +1,77 @@ +/* + * Licensed 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 com.facebook.presto.client.auth.external; + +import com.facebook.presto.client.ClientException; +import com.google.common.annotations.VisibleForTesting; + +import java.net.URI; +import java.time.Duration; +import java.util.Optional; + +import static java.util.Objects.requireNonNull; + +class ExternalAuthentication +{ + private final URI tokenUri; + private final Optional redirectUri; + + public ExternalAuthentication(URI tokenUri, Optional redirectUri) + { + this.tokenUri = requireNonNull(tokenUri, "tokenUri is null"); + this.redirectUri = requireNonNull(redirectUri, "redirectUri is null"); + } + + public Optional obtainToken(Duration timeout, RedirectHandler handler, TokenPoller poller) + { + redirectUri.ifPresent(handler::redirectTo); + + URI currentUri = tokenUri; + + long start = System.nanoTime(); + long timeoutNanos = timeout.toNanos(); + + while (true) { + long remaining = timeoutNanos - (System.nanoTime() - start); + if (remaining < 0) { + return Optional.empty(); + } + + TokenPollResult result = poller.pollForToken(currentUri, Duration.ofNanos(remaining)); + + if (result.isFailed()) { + throw new ClientException(result.getError()); + } + + if (result.isPending()) { + currentUri = result.getNextTokenUri(); + continue; + } + poller.tokenReceived(currentUri); + return Optional.of(result.getToken()); + } + } + + @VisibleForTesting + Optional getRedirectUri() + { + return redirectUri; + } + + @VisibleForTesting + URI getTokenUri() + { + return tokenUri; + } +} diff --git a/presto-client/src/main/java/com/facebook/presto/client/auth/external/ExternalAuthenticator.java b/presto-client/src/main/java/com/facebook/presto/client/auth/external/ExternalAuthenticator.java new file mode 100644 index 0000000000000..60f568dc7c173 --- /dev/null +++ b/presto-client/src/main/java/com/facebook/presto/client/auth/external/ExternalAuthenticator.java @@ -0,0 +1,120 @@ +/* + * Licensed 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 com.facebook.presto.client.auth.external; + +import com.facebook.presto.client.ClientException; +import com.google.common.annotations.VisibleForTesting; +import jakarta.annotation.Nullable; +import okhttp3.Authenticator; +import okhttp3.Challenge; +import okhttp3.Interceptor; +import okhttp3.Request; +import okhttp3.Response; +import okhttp3.Route; + +import java.io.IOException; +import java.net.URI; +import java.net.URISyntaxException; +import java.time.Duration; +import java.util.Map; +import java.util.Optional; + +import static com.google.common.net.HttpHeaders.AUTHORIZATION; +import static java.lang.String.format; +import static java.util.Objects.requireNonNull; + +public class ExternalAuthenticator + implements Authenticator, Interceptor +{ + public static final String TOKEN_URI_FIELD = "x_token_server"; + public static final String REDIRECT_URI_FIELD = "x_redirect_server"; + + private final TokenPoller tokenPoller; + private final RedirectHandler redirectHandler; + private final Duration timeout; + private final KnownToken knownToken; + + public ExternalAuthenticator(RedirectHandler redirect, TokenPoller tokenPoller, KnownToken knownToken, Duration timeout) + { + this.tokenPoller = requireNonNull(tokenPoller, "tokenPoller is null"); + this.redirectHandler = requireNonNull(redirect, "redirect is null"); + this.knownToken = requireNonNull(knownToken, "knownToken is null"); + this.timeout = requireNonNull(timeout, "timeout is null"); + } + + @Nullable + @Override + public Request authenticate(Route route, Response response) + { + knownToken.setupToken(() -> { + Optional authentication = toAuthentication(response); + if (!authentication.isPresent()) { + return Optional.empty(); + } + + return authentication.get().obtainToken(timeout, redirectHandler, tokenPoller); + }); + + return knownToken.getToken() + .map(token -> withBearerToken(response.request(), token)) + .orElse(null); + } + + @Override + public Response intercept(Chain chain) + throws IOException + { + Optional token = knownToken.getToken(); + if (token.isPresent()) { + return chain.proceed(withBearerToken(chain.request(), token.get())); + } + + return chain.proceed(chain.request()); + } + + private static Request withBearerToken(Request request, Token token) + { + return request.newBuilder() + .header(AUTHORIZATION, "Bearer " + token.token()) + .build(); + } + + @VisibleForTesting + static Optional toAuthentication(Response response) + { + for (Challenge challenge : response.challenges()) { + if (challenge.scheme().equalsIgnoreCase("Bearer")) { + Optional tokenUri = parseField(challenge.authParams(), TOKEN_URI_FIELD); + Optional redirectUri = parseField(challenge.authParams(), REDIRECT_URI_FIELD); + if (tokenUri.isPresent()) { + return Optional.of(new ExternalAuthentication(tokenUri.get(), redirectUri)); + } + } + } + + return Optional.empty(); + } + + private static Optional parseField(Map fields, String key) + { + return Optional.ofNullable(fields.get(key)).map(value -> { + try { + return new URI(value); + } + catch (URISyntaxException e) { + throw new ClientException(format("Failed to parse URI for field '%s'", key), e); + } + }); + } +} diff --git a/presto-client/src/main/java/com/facebook/presto/client/auth/external/ExternalRedirectStrategy.java b/presto-client/src/main/java/com/facebook/presto/client/auth/external/ExternalRedirectStrategy.java new file mode 100644 index 0000000000000..07b004f0c15fe --- /dev/null +++ b/presto-client/src/main/java/com/facebook/presto/client/auth/external/ExternalRedirectStrategy.java @@ -0,0 +1,40 @@ +/* + * Licensed 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 com.facebook.presto.client.auth.external; + +import com.google.common.collect.ImmutableList; + +import static java.util.Objects.requireNonNull; + +public enum ExternalRedirectStrategy +{ + DESKTOP_OPEN(new DesktopBrowserRedirectHandler()), + SYSTEM_OPEN(new SystemOpenRedirectHandler()), + PRINT(new SystemOutPrintRedirectHandler()), + OPEN(new CompositeRedirectHandler(ImmutableList.of(SYSTEM_OPEN, DESKTOP_OPEN))), + ALL(new CompositeRedirectHandler(ImmutableList.of(OPEN, PRINT))) + /**/; + + private final RedirectHandler handler; + + ExternalRedirectStrategy(RedirectHandler handler) + { + this.handler = requireNonNull(handler, "handler is null"); + } + + public RedirectHandler getHandler() + { + return handler; + } +} diff --git a/presto-client/src/main/java/com/facebook/presto/client/auth/external/HttpTokenPoller.java b/presto-client/src/main/java/com/facebook/presto/client/auth/external/HttpTokenPoller.java new file mode 100644 index 0000000000000..b06b224ee80ac --- /dev/null +++ b/presto-client/src/main/java/com/facebook/presto/client/auth/external/HttpTokenPoller.java @@ -0,0 +1,193 @@ +/* + * Licensed 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 com.facebook.presto.client.auth.external; + +import com.facebook.airlift.json.JsonCodec; +import com.facebook.presto.client.JsonResponse; +import com.fasterxml.jackson.annotation.JsonCreator; +import com.fasterxml.jackson.annotation.JsonProperty; +import net.jodah.failsafe.Failsafe; +import net.jodah.failsafe.FailsafeException; +import net.jodah.failsafe.RetryPolicy; +import okhttp3.HttpUrl; +import okhttp3.OkHttpClient; +import okhttp3.Request; +import okhttp3.Response; + +import java.io.IOException; +import java.io.UncheckedIOException; +import java.net.URI; +import java.time.Duration; +import java.util.Optional; +import java.util.function.Consumer; +import java.util.function.Supplier; + +import static com.facebook.airlift.json.JsonCodec.jsonCodec; +import static com.facebook.presto.client.JsonResponse.execute; +import static com.google.common.base.MoreObjects.firstNonNull; +import static com.google.common.net.HttpHeaders.USER_AGENT; +import static java.lang.String.format; +import static java.net.HttpURLConnection.HTTP_INTERNAL_ERROR; +import static java.net.HttpURLConnection.HTTP_OK; +import static java.net.HttpURLConnection.HTTP_UNAVAILABLE; +import static java.time.temporal.ChronoUnit.MILLIS; +import static java.util.Objects.requireNonNull; + +public class HttpTokenPoller + implements TokenPoller +{ + private static final JsonCodec TOKEN_POLL_CODEC = jsonCodec(TokenPollRepresentation.class); + private static final String USER_AGENT_VALUE = "PrestoTokenPoller/" + + firstNonNull(HttpTokenPoller.class.getPackage().getImplementationVersion(), "unknown"); + + private final Supplier client; + + public HttpTokenPoller(OkHttpClient client) + { + requireNonNull(client, "client is null"); + this.client = () -> client; + } + + public HttpTokenPoller(OkHttpClient client, Consumer refreshableClientConfig) + { + requireNonNull(client, "client is null"); + requireNonNull(refreshableClientConfig, "refreshableClientConfig is null"); + + this.client = () -> { + OkHttpClient.Builder builder = client.newBuilder(); + refreshableClientConfig.accept(builder); + return builder.build(); + }; + } + + @Override + public TokenPollResult pollForToken(URI tokenUri, Duration timeout) + { + try { + return Failsafe.with(new RetryPolicy() + .withMaxAttempts(-1) + .withMaxDuration(timeout) + .withBackoff(100, 500, MILLIS) + .handle(IOException.class)) + .get(() -> executePoll(prepareRequestBuilder(tokenUri).build())); + } + catch (FailsafeException e) { + if (e.getCause() instanceof IOException) { + throw new UncheckedIOException((IOException) e.getCause()); + } + throw e; + } + } + + @Override + public void tokenReceived(URI tokenUri) + { + try { + Failsafe.with(new RetryPolicy() + .withMaxAttempts(-1) + .withMaxDuration(Duration.ofSeconds(4)) + .withBackoff(100, 500, MILLIS) + .handleResultIf(code -> code >= HTTP_INTERNAL_ERROR)) + .get(() -> { + Request request = prepareRequestBuilder(tokenUri) + .delete() + .build(); + try (Response response = client.get().newCall(request) + .execute()) { + return response.code(); + } + }); + } + catch (FailsafeException e) { + if (e.getCause() instanceof IOException) { + throw new UncheckedIOException((IOException) e.getCause()); + } + throw e; + } + } + + private static Request.Builder prepareRequestBuilder(URI tokenUri) + { + HttpUrl url = HttpUrl.get(tokenUri); + if (url == null) { + throw new IllegalArgumentException("Invalid token URI: " + tokenUri); + } + + return new Request.Builder() + .url(url) + .addHeader(USER_AGENT, USER_AGENT_VALUE); + } + + private TokenPollResult executePoll(Request request) + throws IOException + { + JsonResponse response = executeRequest(request); + + if ((response.getStatusCode() == HTTP_OK) && response.hasValue()) { + return response.getValue().toResult(); + } + + Optional responseBody = Optional.ofNullable(response.getResponseBody()); + String message = format("Request to %s failed: %s [Error: %s]", request.url(), response, responseBody.orElse("")); + + if (response.getStatusCode() == HTTP_UNAVAILABLE) { + throw new IOException(message); + } + + return TokenPollResult.failed(message); + } + + private JsonResponse executeRequest(Request request) + throws IOException + { + try { + return execute(TOKEN_POLL_CODEC, client.get(), request); + } + catch (UncheckedIOException e) { + throw e.getCause(); + } + } + + public static class TokenPollRepresentation + { + private final String token; + private final URI nextUri; + private final String error; + + @JsonCreator + public TokenPollRepresentation( + @JsonProperty("token") String token, + @JsonProperty("nextUri") URI nextUri, + @JsonProperty("error") String error) + { + this.token = token; + this.nextUri = nextUri; + this.error = error; + } + + TokenPollResult toResult() + { + if (token != null) { + return TokenPollResult.successful(new Token(token)); + } + if (error != null) { + return TokenPollResult.failed(error); + } + if (nextUri != null) { + return TokenPollResult.pending(nextUri); + } + return TokenPollResult.failed("Failed to poll for token. No fields set in response."); + } + } +} diff --git a/presto-client/src/main/java/com/facebook/presto/client/auth/external/KnownToken.java b/presto-client/src/main/java/com/facebook/presto/client/auth/external/KnownToken.java new file mode 100644 index 0000000000000..ae57de609c02c --- /dev/null +++ b/presto-client/src/main/java/com/facebook/presto/client/auth/external/KnownToken.java @@ -0,0 +1,34 @@ +/* + * Licensed 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 com.facebook.presto.client.auth.external; + +import java.util.Optional; +import java.util.function.Supplier; + +public interface KnownToken +{ + Optional getToken(); + + void setupToken(Supplier> tokenSource); + + static KnownToken local() + { + return new LocalKnownToken(); + } + + static KnownToken memoryCached() + { + return MemoryCachedKnownToken.INSTANCE; + } +} diff --git a/presto-client/src/main/java/com/facebook/presto/client/auth/external/LocalKnownToken.java b/presto-client/src/main/java/com/facebook/presto/client/auth/external/LocalKnownToken.java new file mode 100644 index 0000000000000..31ff4bd89d525 --- /dev/null +++ b/presto-client/src/main/java/com/facebook/presto/client/auth/external/LocalKnownToken.java @@ -0,0 +1,46 @@ +/* + * Licensed 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 com.facebook.presto.client.auth.external; + +import javax.annotation.concurrent.NotThreadSafe; + +import java.util.Optional; +import java.util.function.Supplier; + +import static java.util.Objects.requireNonNull; + +/** + * LocalKnownToken class keeps the token on its field + * and it's designed to use it in fully serialized manner. + */ +@NotThreadSafe +class LocalKnownToken + implements KnownToken +{ + private Optional knownToken = Optional.empty(); + + @Override + public Optional getToken() + { + return knownToken; + } + + @Override + public void setupToken(Supplier> tokenSource) + { + requireNonNull(tokenSource, "tokenSource is null"); + + knownToken = tokenSource.get(); + } +} diff --git a/presto-client/src/main/java/com/facebook/presto/client/auth/external/MemoryCachedKnownToken.java b/presto-client/src/main/java/com/facebook/presto/client/auth/external/MemoryCachedKnownToken.java new file mode 100644 index 0000000000000..0b1253a6deaa6 --- /dev/null +++ b/presto-client/src/main/java/com/facebook/presto/client/auth/external/MemoryCachedKnownToken.java @@ -0,0 +1,83 @@ +/* + * Licensed 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 com.facebook.presto.client.auth.external; + +import javax.annotation.concurrent.ThreadSafe; + +import java.util.Optional; +import java.util.concurrent.locks.Lock; +import java.util.concurrent.locks.ReadWriteLock; +import java.util.concurrent.locks.ReentrantReadWriteLock; +import java.util.function.Supplier; + +/** + * This KnownToken instance forces all Connections to reuse same token. + * Every time an existing token is considered to be invalid each Connection + * will try to obtain a new token, but only the first one will actually do the job, + * where every other connection will be waiting on readLock + * until obtaining new token finishes. + *

+ * In general the game is to reuse same token and obtain it only once, no matter how + * many Connections will be actively using it. It's very important as obtaining the new token + * will take minutes, as it mostly requires user thinking time. + */ +@ThreadSafe +public class MemoryCachedKnownToken + implements KnownToken +{ + public static final MemoryCachedKnownToken INSTANCE = new MemoryCachedKnownToken(); + + private final ReadWriteLock lock = new ReentrantReadWriteLock(); + private final Lock readLock = lock.readLock(); + private final Lock writeLock = lock.writeLock(); + private Optional knownToken = Optional.empty(); + + private MemoryCachedKnownToken() + { + } + + @Override + public Optional getToken() + { + try { + readLock.lockInterruptibly(); + return knownToken; + } + catch (InterruptedException e) { + Thread.currentThread().interrupt(); + throw new RuntimeException(e); + } + finally { + readLock.unlock(); + } + } + + @Override + public void setupToken(Supplier> tokenSource) + { + // Try to lock and generate new token. If some other thread (Connection) has + // already obtained writeLock and is generating new token, then skip this + // to block on getToken() + if (writeLock.tryLock()) { + try { + // Clear knownToken before obtaining new token, as it might fail leaving old invalid token. + knownToken = Optional.empty(); + knownToken = tokenSource.get(); + } + finally { + writeLock.unlock(); + } + } + } +} diff --git a/presto-client/src/main/java/com/facebook/presto/client/auth/external/RedirectException.java b/presto-client/src/main/java/com/facebook/presto/client/auth/external/RedirectException.java new file mode 100644 index 0000000000000..0c5a1f5e149fa --- /dev/null +++ b/presto-client/src/main/java/com/facebook/presto/client/auth/external/RedirectException.java @@ -0,0 +1,28 @@ +/* + * Licensed 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 com.facebook.presto.client.auth.external; + +public class RedirectException + extends RuntimeException +{ + public RedirectException(String message, Throwable throwable) + { + super(message, throwable); + } + + public RedirectException(String message) + { + super(message); + } +} diff --git a/presto-client/src/main/java/com/facebook/presto/client/auth/external/RedirectHandler.java b/presto-client/src/main/java/com/facebook/presto/client/auth/external/RedirectHandler.java new file mode 100644 index 0000000000000..710be10238884 --- /dev/null +++ b/presto-client/src/main/java/com/facebook/presto/client/auth/external/RedirectHandler.java @@ -0,0 +1,22 @@ +/* + * Licensed 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 com.facebook.presto.client.auth.external; + +import java.net.URI; + +public interface RedirectHandler +{ + void redirectTo(URI uri) + throws RedirectException; +} diff --git a/presto-client/src/main/java/com/facebook/presto/client/auth/external/SystemOpenRedirectHandler.java b/presto-client/src/main/java/com/facebook/presto/client/auth/external/SystemOpenRedirectHandler.java new file mode 100644 index 0000000000000..d63841b26d830 --- /dev/null +++ b/presto-client/src/main/java/com/facebook/presto/client/auth/external/SystemOpenRedirectHandler.java @@ -0,0 +1,102 @@ +/* + * Licensed 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 com.facebook.presto.client.auth.external; + +import com.google.common.base.Splitter; +import com.google.common.collect.ImmutableList; + +import java.io.File; +import java.io.FileNotFoundException; +import java.io.IOException; +import java.net.URI; +import java.nio.file.Paths; +import java.util.List; +import java.util.Optional; + +import static java.lang.String.format; +import static java.util.Locale.ENGLISH; + +public class SystemOpenRedirectHandler + implements RedirectHandler +{ + private static final List LINUX_BROWSERS = ImmutableList.of( + "xdg-open", + "gnome-open", + "kde-open", + "chromium", + "google", + "google-chrome", + "firefox", + "mozilla", + "opera", + "epiphany", + "konqueror"); + + private static final String MACOS_OPEN_COMMAND = "open"; + private static final String WINDOWS_OPEN_COMMAND = "rundll32 url.dll,FileProtocolHandler"; + + private static final Splitter SPLITTER = Splitter.on(":") + .omitEmptyStrings() + .trimResults(); + + @Override + public void redirectTo(URI uri) + throws RedirectException + { + String operatingSystem = System.getProperty("os.name").toLowerCase(ENGLISH); + + try { + if (operatingSystem.contains("mac")) { + exec(uri, MACOS_OPEN_COMMAND); + } + else if (operatingSystem.contains("windows")) { + exec(uri, WINDOWS_OPEN_COMMAND); + } + else { + String executablePath = findLinuxBrowser() + .orElseThrow(() -> new FileNotFoundException("Could not find any known linux browser in $PATH")); + exec(uri, executablePath); + } + } + catch (IOException e) { + throw new RedirectException(format("Could not open uri %s", uri), e); + } + } + + private static Optional findLinuxBrowser() + { + List paths = SPLITTER.splitToList(System.getenv("PATH")); + for (String path : paths) { + File[] found = Paths.get(path) + .toFile() + .listFiles((dir, name) -> LINUX_BROWSERS.contains(name)); + + if (found == null) { + continue; + } + + if (found.length > 0) { + return Optional.of(found[0].getPath()); + } + } + + return Optional.empty(); + } + + private static void exec(URI uri, String openCommand) + throws IOException + { + Runtime.getRuntime().exec(openCommand + " " + uri.toString()); + } +} diff --git a/presto-client/src/main/java/com/facebook/presto/client/auth/external/SystemOutPrintRedirectHandler.java b/presto-client/src/main/java/com/facebook/presto/client/auth/external/SystemOutPrintRedirectHandler.java new file mode 100644 index 0000000000000..d0f3c7954dbd7 --- /dev/null +++ b/presto-client/src/main/java/com/facebook/presto/client/auth/external/SystemOutPrintRedirectHandler.java @@ -0,0 +1,27 @@ +/* + * Licensed 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 com.facebook.presto.client.auth.external; + +import java.net.URI; + +public class SystemOutPrintRedirectHandler + implements RedirectHandler +{ + @Override + public void redirectTo(URI uri) throws RedirectException + { + System.out.println("External authentication required. Please go to:"); + System.out.println(uri.toString()); + } +} diff --git a/presto-client/src/main/java/com/facebook/presto/client/auth/external/Token.java b/presto-client/src/main/java/com/facebook/presto/client/auth/external/Token.java new file mode 100644 index 0000000000000..82890f3f6b2df --- /dev/null +++ b/presto-client/src/main/java/com/facebook/presto/client/auth/external/Token.java @@ -0,0 +1,31 @@ +/* + * Licensed 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 com.facebook.presto.client.auth.external; + +import static java.util.Objects.requireNonNull; + +class Token +{ + private final String token; + + public Token(String token) + { + this.token = requireNonNull(token, "token is null"); + } + + public String token() + { + return token; + } +} diff --git a/presto-client/src/main/java/com/facebook/presto/client/auth/external/TokenPollResult.java b/presto-client/src/main/java/com/facebook/presto/client/auth/external/TokenPollResult.java new file mode 100644 index 0000000000000..3cd13b5dd3710 --- /dev/null +++ b/presto-client/src/main/java/com/facebook/presto/client/auth/external/TokenPollResult.java @@ -0,0 +1,80 @@ +/* + * Licensed 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 com.facebook.presto.client.auth.external; + +import java.net.URI; +import java.util.Optional; + +import static java.util.Objects.requireNonNull; + +public class TokenPollResult +{ + private enum State + { + PENDING, SUCCESSFUL, FAILED + } + + private final State state; + private final Optional errorMessage; + private final Optional nextTokenUri; + private final Optional token; + + public static TokenPollResult failed(String error) + { + return new TokenPollResult(State.FAILED, null, error, null); + } + + public static TokenPollResult pending(URI uri) + { + return new TokenPollResult(State.PENDING, null, null, uri); + } + + public static TokenPollResult successful(Token token) + { + return new TokenPollResult(State.SUCCESSFUL, token, null, null); + } + + private TokenPollResult(State state, Token token, String error, URI nextTokenUri) + { + this.state = requireNonNull(state, "state is null"); + this.token = Optional.ofNullable(token); + this.errorMessage = Optional.ofNullable(error); + this.nextTokenUri = Optional.ofNullable(nextTokenUri); + } + + public boolean isFailed() + { + return state == State.FAILED; + } + + public boolean isPending() + { + return state == State.PENDING; + } + + public Token getToken() + { + return token.orElseThrow(() -> new IllegalStateException("state is " + state)); + } + + public String getError() + { + return errorMessage.orElseThrow(() -> new IllegalStateException("state is " + state)); + } + + public URI getNextTokenUri() + { + return nextTokenUri.orElseThrow(() -> new IllegalStateException("state is " + state)); + } +} diff --git a/presto-client/src/main/java/com/facebook/presto/client/auth/external/TokenPoller.java b/presto-client/src/main/java/com/facebook/presto/client/auth/external/TokenPoller.java new file mode 100644 index 0000000000000..ea228137ab688 --- /dev/null +++ b/presto-client/src/main/java/com/facebook/presto/client/auth/external/TokenPoller.java @@ -0,0 +1,24 @@ +/* + * Licensed 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 com.facebook.presto.client.auth.external; + +import java.net.URI; +import java.time.Duration; + +public interface TokenPoller +{ + TokenPollResult pollForToken(URI tokenUri, Duration timeout); + + void tokenReceived(URI tokenUri); +} diff --git a/presto-client/src/test/java/com/facebook/presto/client/auth/external/MockRedirectHandler.java b/presto-client/src/test/java/com/facebook/presto/client/auth/external/MockRedirectHandler.java new file mode 100644 index 0000000000000..f4b8c6cd07633 --- /dev/null +++ b/presto-client/src/test/java/com/facebook/presto/client/auth/external/MockRedirectHandler.java @@ -0,0 +1,59 @@ +/* + * Licensed 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 com.facebook.presto.client.auth.external; + +import java.net.URI; +import java.time.Duration; +import java.util.concurrent.atomic.AtomicInteger; + +public class MockRedirectHandler + implements RedirectHandler +{ + private URI redirectedTo; + private AtomicInteger redirectionCount = new AtomicInteger(0); + private Duration redirectTime; + + @Override + public void redirectTo(URI uri) + throws RedirectException + { + redirectedTo = uri; + redirectionCount.incrementAndGet(); + try { + if (redirectTime != null) { + Thread.sleep(redirectTime.toMillis()); + } + } + catch (InterruptedException e) { + Thread.currentThread().interrupt(); + throw new RuntimeException(e); + } + } + + public URI redirectedTo() + { + return redirectedTo; + } + + public int getRedirectionCount() + { + return redirectionCount.get(); + } + + public MockRedirectHandler sleepOnRedirect(Duration redirectTime) + { + this.redirectTime = redirectTime; + return this; + } +} diff --git a/presto-client/src/test/java/com/facebook/presto/client/auth/external/MockTokenPoller.java b/presto-client/src/test/java/com/facebook/presto/client/auth/external/MockTokenPoller.java new file mode 100644 index 0000000000000..6733a76ebb64d --- /dev/null +++ b/presto-client/src/test/java/com/facebook/presto/client/auth/external/MockTokenPoller.java @@ -0,0 +1,81 @@ +/* + * Licensed 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 com.facebook.presto.client.auth.external; + +import com.google.common.collect.ImmutableList; + +import java.net.URI; +import java.time.Duration; +import java.util.Map; +import java.util.concurrent.BlockingDeque; +import java.util.concurrent.ConcurrentHashMap; +import java.util.concurrent.LinkedBlockingDeque; +import java.util.function.Function; + +public final class MockTokenPoller + implements TokenPoller +{ + private final Map> results = new ConcurrentHashMap<>(); + private URI tokenReceivedUri; + + public static TokenPoller onPoll(Function pollingStrategy) + { + return new TokenPoller() + { + @Override + public TokenPollResult pollForToken(URI tokenUri, Duration timeout) + { + return pollingStrategy.apply(tokenUri); + } + + @Override + public void tokenReceived(URI tokenUri) + { + } + }; + } + + public MockTokenPoller withResult(URI tokenUri, TokenPollResult result) + { + results.compute(tokenUri, (uri, queue) -> { + if (queue == null) { + return new LinkedBlockingDeque<>(ImmutableList.of(result)); + } + queue.add(result); + return queue; + }); + return this; + } + + @Override + public TokenPollResult pollForToken(URI tokenUri, Duration ignored) + { + BlockingDeque queue = results.get(tokenUri); + if (queue == null) { + throw new IllegalArgumentException("Unknown token URI: " + tokenUri); + } + return queue.remove(); + } + + @Override + public void tokenReceived(URI tokenUri) + { + this.tokenReceivedUri = tokenUri; + } + + public URI tokenReceivedUri() + { + return tokenReceivedUri; + } +} diff --git a/presto-client/src/test/java/com/facebook/presto/client/auth/external/TestExternalAuthentication.java b/presto-client/src/test/java/com/facebook/presto/client/auth/external/TestExternalAuthentication.java new file mode 100644 index 0000000000000..e7434c5ac242e --- /dev/null +++ b/presto-client/src/test/java/com/facebook/presto/client/auth/external/TestExternalAuthentication.java @@ -0,0 +1,130 @@ +/* + * Licensed 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 com.facebook.presto.client.auth.external; + +import com.facebook.presto.client.ClientException; +import org.testng.annotations.Test; + +import java.io.IOException; +import java.io.UncheckedIOException; +import java.net.URI; +import java.time.Duration; +import java.util.Optional; +import java.util.concurrent.TimeUnit; + +import static com.google.common.util.concurrent.Uninterruptibles.sleepUninterruptibly; +import static java.net.URI.create; +import static org.assertj.core.api.Assertions.assertThat; +import static org.assertj.core.api.Assertions.assertThatThrownBy; + +public class TestExternalAuthentication +{ + private static final String AUTH_TOKEN = "authToken"; + private static final URI REDIRECT_URI = create("https://redirect.uri"); + private static final URI TOKEN_URI = create("https://token.uri"); + private static final Duration TIMEOUT = Duration.ofSeconds(1); + + @Test + public void testObtainTokenWhenTokenAlreadyExists() + { + MockRedirectHandler redirectHandler = new MockRedirectHandler(); + + MockTokenPoller poller = new MockTokenPoller() + .withResult(TOKEN_URI, TokenPollResult.successful(new Token(AUTH_TOKEN))); + + Optional token = new ExternalAuthentication(TOKEN_URI, Optional.of(REDIRECT_URI)) + .obtainToken(TIMEOUT, redirectHandler, poller); + + assertThat(redirectHandler.redirectedTo()).isEqualTo(REDIRECT_URI); + assertThat(token).map(Token::token).hasValue(AUTH_TOKEN); + assertThat(poller.tokenReceivedUri()).isEqualTo(TOKEN_URI); + } + + @Test + public void testObtainTokenWhenTokenIsReadyAtSecondAttempt() + { + RedirectHandler redirectHandler = new MockRedirectHandler(); + + URI nextTokenUri = TOKEN_URI.resolve("/next"); + MockTokenPoller poller = new MockTokenPoller() + .withResult(TOKEN_URI, TokenPollResult.pending(nextTokenUri)) + .withResult(nextTokenUri, TokenPollResult.successful(new Token(AUTH_TOKEN))); + + Optional token = new ExternalAuthentication(TOKEN_URI, Optional.of(REDIRECT_URI)) + .obtainToken(TIMEOUT, redirectHandler, poller); + + assertThat(token).map(Token::token).hasValue(AUTH_TOKEN); + assertThat(poller.tokenReceivedUri()).isEqualTo(nextTokenUri); + } + + @Test + public void testObtainTokenWhenTokenIsNeverAvailable() + { + RedirectHandler redirectHandler = new MockRedirectHandler(); + + TokenPoller poller = MockTokenPoller.onPoll(tokenUri -> { + sleepUninterruptibly(20, TimeUnit.MILLISECONDS); + return TokenPollResult.pending(TOKEN_URI); + }); + + Optional token = new ExternalAuthentication(TOKEN_URI, Optional.of(REDIRECT_URI)) + .obtainToken(TIMEOUT, redirectHandler, poller); + + assertThat(token).isEmpty(); + } + + @Test + public void testObtainTokenWhenPollingFails() + { + RedirectHandler redirectHandler = new MockRedirectHandler(); + + TokenPoller poller = new MockTokenPoller() + .withResult(TOKEN_URI, TokenPollResult.failed("error")); + + assertThatThrownBy(() -> new ExternalAuthentication(TOKEN_URI, Optional.of(REDIRECT_URI)) + .obtainToken(TIMEOUT, redirectHandler, poller)) + .isInstanceOf(ClientException.class) + .hasMessage("error"); + } + + @Test + public void testObtainTokenWhenPollingFailsWithException() + { + RedirectHandler redirectHandler = new MockRedirectHandler(); + + TokenPoller poller = MockTokenPoller.onPoll(uri -> { + throw new UncheckedIOException(new IOException("polling error")); + }); + + assertThatThrownBy(() -> new ExternalAuthentication(TOKEN_URI, Optional.of(REDIRECT_URI)) + .obtainToken(TIMEOUT, redirectHandler, poller)) + .isInstanceOf(UncheckedIOException.class) + .hasRootCauseInstanceOf(IOException.class); + } + + @Test + public void testObtainTokenWhenNoRedirectUriHasBeenProvided() + { + MockRedirectHandler redirectHandler = new MockRedirectHandler(); + + TokenPoller poller = new MockTokenPoller() + .withResult(TOKEN_URI, TokenPollResult.successful(new Token(AUTH_TOKEN))); + + Optional token = new ExternalAuthentication(TOKEN_URI, Optional.empty()) + .obtainToken(TIMEOUT, redirectHandler, poller); + + assertThat(redirectHandler.redirectedTo()).isNull(); + assertThat(token).map(Token::token).hasValue(AUTH_TOKEN); + } +} diff --git a/presto-client/src/test/java/com/facebook/presto/client/auth/external/TestExternalAuthenticator.java b/presto-client/src/test/java/com/facebook/presto/client/auth/external/TestExternalAuthenticator.java new file mode 100644 index 0000000000000..86cd89fc6dfa9 --- /dev/null +++ b/presto-client/src/test/java/com/facebook/presto/client/auth/external/TestExternalAuthenticator.java @@ -0,0 +1,370 @@ +/* + * Licensed 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 com.facebook.presto.client.auth.external; + +import com.facebook.presto.client.ClientException; +import com.google.common.collect.ImmutableList; +import okhttp3.HttpUrl; +import okhttp3.Protocol; +import okhttp3.Request; +import okhttp3.Response; +import org.assertj.core.api.ListAssert; +import org.assertj.core.api.ThrowableAssert; +import org.testng.annotations.AfterClass; +import org.testng.annotations.Test; + +import java.net.URI; +import java.net.URISyntaxException; +import java.time.Duration; +import java.util.ArrayList; +import java.util.List; +import java.util.Optional; +import java.util.concurrent.Callable; +import java.util.concurrent.CancellationException; +import java.util.concurrent.ExecutionException; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Future; +import java.util.stream.Stream; + +import static com.facebook.airlift.concurrent.Threads.daemonThreadsNamed; +import static com.facebook.presto.client.auth.external.ExternalAuthenticator.TOKEN_URI_FIELD; +import static com.facebook.presto.client.auth.external.ExternalAuthenticator.toAuthentication; +import static com.facebook.presto.client.auth.external.MockTokenPoller.onPoll; +import static com.facebook.presto.client.auth.external.TokenPollResult.successful; +import static com.google.common.base.Preconditions.checkState; +import static com.google.common.collect.ImmutableList.toImmutableList; +import static com.google.common.net.HttpHeaders.AUTHORIZATION; +import static com.google.common.net.HttpHeaders.WWW_AUTHENTICATE; +import static java.lang.String.format; +import static java.net.HttpURLConnection.HTTP_UNAUTHORIZED; +import static java.net.URI.create; +import static java.util.concurrent.Executors.newCachedThreadPool; +import static org.assertj.core.api.Assertions.assertThat; +import static org.assertj.core.api.Assertions.assertThatThrownBy; + +@Test(singleThreaded = true) +public class TestExternalAuthenticator +{ + private static final ExecutorService executor = newCachedThreadPool(daemonThreadsNamed(TestExternalAuthenticator.class.getName() + "-%d")); + + @AfterClass(alwaysRun = true) + public void shutDownThreadPool() + { + executor.shutdownNow(); + } + + @Test + public void testChallengeWithOnlyTokenServerUri() + { + assertThat(buildAuthentication("Bearer x_token_server=\"http://token.uri\"")) + .hasValueSatisfying(authentication -> { + assertThat(authentication.getRedirectUri()).isEmpty(); + assertThat(authentication.getTokenUri()).isEqualTo(create("http://token.uri")); + }); + } + + @Test + public void testChallengeWithBothUri() + { + assertThat(buildAuthentication("Bearer x_redirect_server=\"http://redirect.uri\", x_token_server=\"http://token.uri\"")) + .hasValueSatisfying(authentication -> { + assertThat(authentication.getRedirectUri()).hasValue(create("http://redirect.uri")); + assertThat(authentication.getTokenUri()).isEqualTo(create("http://token.uri")); + }); + } + + @Test + public void testChallengeWithValuesWithoutQuotes() + { + // this is legal according to RFC 7235 + assertThat(buildAuthentication("Bearer x_redirect_server=http://redirect.uri, x_token_server=http://token.uri")) + .hasValueSatisfying(authentication -> { + assertThat(authentication.getRedirectUri()).hasValue(create("http://redirect.uri")); + assertThat(authentication.getTokenUri()).isEqualTo(create("http://token.uri")); + }); + } + + @Test + public void testChallengeWithAdditionalFields() + { + assertThat(buildAuthentication("Bearer type=\"token\", x_redirect_server=\"http://redirect.uri\", x_token_server=\"http://token.uri\", description=\"oauth challenge\"")) + .hasValueSatisfying(authentication -> { + assertThat(authentication.getRedirectUri()).hasValue(create("http://redirect.uri")); + assertThat(authentication.getTokenUri()).isEqualTo(create("http://token.uri")); + }); + } + + @Test + public void testInvalidChallenges() + { + // no authentication parameters + assertThat(buildAuthentication("Bearer")).isEmpty(); + + // no Bearer scheme prefix + assertThat(buildAuthentication("x_redirect_server=\"http://redirect.uri\", x_token_server=\"http://token.uri\"")).isEmpty(); + + // space instead of comma + assertThat(buildAuthentication("Bearer x_redirect_server=\"http://redirect.uri\" x_token_server=\"http://token.uri\"")).isEmpty(); + + // equals sign instead of comma + assertThat(buildAuthentication("Bearer x_redirect_server=\"http://redirect.uri\"=x_token_server=\"http://token.uri\"")).isEmpty(); + } + + @Test + public void testChallengeWithMalformedUri() + { + assertThatThrownBy(() -> buildAuthentication("Bearer x_token_server=\"http://[1.1.1.1]\"")) + .isInstanceOf(ClientException.class) + .hasMessageContaining(format("Failed to parse URI for field '%s'", TOKEN_URI_FIELD)) + .hasRootCauseInstanceOf(URISyntaxException.class); + } + + @Test + public void testAuthentication() + { + MockTokenPoller tokenPoller = new MockTokenPoller() + .withResult(URI.create("http://token.uri"), successful(new Token("valid-token"))); + ExternalAuthenticator authenticator = new ExternalAuthenticator(uri -> {}, tokenPoller, KnownToken.local(), Duration.ofSeconds(1)); + + Request authenticated = authenticator.authenticate(null, getUnauthorizedResponse("Bearer x_token_server=\"http://token.uri\"")); + + assertThat(authenticated.headers(AUTHORIZATION)) + .containsExactly("Bearer valid-token"); + } + + @Test + public void testReAuthenticationAfterRejectingToken() + { + MockTokenPoller tokenPoller = new MockTokenPoller() + .withResult(URI.create("http://token.uri"), successful(new Token("first-token"))) + .withResult(URI.create("http://token.uri"), successful(new Token("second-token"))); + ExternalAuthenticator authenticator = new ExternalAuthenticator(uri -> {}, tokenPoller, KnownToken.local(), Duration.ofSeconds(1)); + + Request request = authenticator.authenticate(null, getUnauthorizedResponse("Bearer x_token_server=\"http://token.uri\"")); + Request reAuthenticated = authenticator.authenticate(null, getUnauthorizedResponse("Bearer x_token_server=\"http://token.uri\"", request)); + + assertThat(reAuthenticated.headers(AUTHORIZATION)) + .containsExactly("Bearer second-token"); + } + + @Test(timeOut = 2000) + public void testAuthenticationFromMultipleThreadsWithLocallyStoredToken() + { + MockTokenPoller tokenPoller = new MockTokenPoller() + .withResult(URI.create("http://token.uri"), successful(new Token("valid-token-1"))) + .withResult(URI.create("http://token.uri"), successful(new Token("valid-token-2"))) + .withResult(URI.create("http://token.uri"), successful(new Token("valid-token-3"))) + .withResult(URI.create("http://token.uri"), successful(new Token("valid-token-4"))); + MockRedirectHandler redirectHandler = new MockRedirectHandler(); + + List> requests = times( + 4, + () -> new ExternalAuthenticator(redirectHandler, tokenPoller, KnownToken.local(), Duration.ofSeconds(1)) + .authenticate(null, getUnauthorizedResponse("Bearer x_token_server=\"http://token.uri\", x_redirect_server=\"http://redirect.uri\""))) + .map(executor::submit) + .collect(toImmutableList()); + + ConcurrentRequestAssertion assertion = new ConcurrentRequestAssertion(requests); + assertion.requests() + .extracting(Request::headers) + .extracting(headers -> headers.get(AUTHORIZATION)) + .contains("Bearer valid-token-1", "Bearer valid-token-2", "Bearer valid-token-3", "Bearer valid-token-4"); + assertion.assertThatNoExceptionsHasBeenThrown(); + assertThat(redirectHandler.getRedirectionCount()).isEqualTo(4); + } + + @Test(timeOut = 2000) + public void testAuthenticationFromMultipleThreadsWithCachedToken() + { + MockTokenPoller tokenPoller = new MockTokenPoller() + .withResult(URI.create("http://token.uri"), successful(new Token("valid-token"))); + MockRedirectHandler redirectHandler = new MockRedirectHandler() + .sleepOnRedirect(Duration.ofSeconds(1)); + + List> requests = times( + 2, + () -> new ExternalAuthenticator(redirectHandler, tokenPoller, KnownToken.memoryCached(), Duration.ofSeconds(1)) + .authenticate(null, getUnauthorizedResponse("Bearer x_token_server=\"http://token.uri\", x_redirect_server=\"http://redirect.uri\""))) + .map(executor::submit) + .collect(toImmutableList()); + + ConcurrentRequestAssertion assertion = new ConcurrentRequestAssertion(requests); + assertion.requests() + .extracting(Request::headers) + .extracting(headers -> headers.get(AUTHORIZATION)) + .containsOnly("Bearer valid-token"); + assertion.assertThatNoExceptionsHasBeenThrown(); + assertThat(redirectHandler.getRedirectionCount()).isEqualTo(1); + } + + @Test(timeOut = 2000) + public void testAuthenticationFromMultipleThreadsWithCachedTokenAfterAuthenticateFails() + { + MockTokenPoller tokenPoller = new MockTokenPoller() + .withResult(URI.create("http://token.uri"), TokenPollResult.successful(new Token("first-token"))) + .withResult(URI.create("http://token.uri"), TokenPollResult.failed("external authentication error")); + MockRedirectHandler redirectHandler = new MockRedirectHandler() + .sleepOnRedirect(Duration.ofMillis(500)); + + ExternalAuthenticator authenticator = new ExternalAuthenticator(redirectHandler, tokenPoller, KnownToken.memoryCached(), Duration.ofSeconds(1)); + Request firstRequest = authenticator.authenticate(null, getUnauthorizedResponse("Bearer x_token_server=\"http://token.uri\", x_redirect_server=\"http://redirect.uri\"")); + + List> requests = times( + 4, + () -> new ExternalAuthenticator(redirectHandler, tokenPoller, KnownToken.memoryCached(), Duration.ofSeconds(1)) + .authenticate(null, getUnauthorizedResponse("Bearer x_token_server=\"http://token.uri\", x_redirect_server=\"http://redirect.uri\"", firstRequest))) + .map(executor::submit) + .collect(toImmutableList()); + + ConcurrentRequestAssertion assertion = new ConcurrentRequestAssertion(requests); + assertion.firstException().hasMessage("external authentication error") + .isInstanceOf(ClientException.class); + + assertThat(redirectHandler.getRedirectionCount()).isEqualTo(2); + } + + @Test(timeOut = 2000) + public void testAuthenticationFromMultipleThreadsWithCachedTokenAfterAuthenticateTimesOut() + { + MockRedirectHandler redirectHandler = new MockRedirectHandler() + .sleepOnRedirect(Duration.ofSeconds(1)); + + List> requests = times( + 2, + () -> new ExternalAuthenticator(redirectHandler, onPoll(TokenPollResult::pending), KnownToken.memoryCached(), Duration.ofMillis(1)) + .authenticate(null, getUnauthorizedResponse("Bearer x_token_server=\"http://token.uri\", x_redirect_server=\"http://redirect.uri\""))) + .map(executor::submit) + .collect(toImmutableList()); + + ConcurrentRequestAssertion assertion = new ConcurrentRequestAssertion(requests); + assertion.requests() + .containsExactly(null, null); + assertion.assertThatNoExceptionsHasBeenThrown(); + assertThat(redirectHandler.getRedirectionCount()).isEqualTo(1); + } + + @Test(timeOut = 2000) + public void testAuthenticationFromMultipleThreadsWithCachedTokenAfterAuthenticateIsInterrupted() + throws Exception + { + ExecutorService interruptableThreadPool = newCachedThreadPool(daemonThreadsNamed(this.getClass().getName() + "-interruptable-%d")); + MockRedirectHandler redirectHandler = new MockRedirectHandler() + .sleepOnRedirect(Duration.ofMinutes(1)); + + ExternalAuthenticator authenticator = new ExternalAuthenticator(redirectHandler, onPoll(TokenPollResult::pending), KnownToken.memoryCached(), Duration.ofMillis(1)); + Future interruptedAuthentication = interruptableThreadPool.submit( + () -> authenticator.authenticate(null, getUnauthorizedResponse("Bearer x_token_server=\"http://token.uri\", x_redirect_server=\"http://redirect.uri\""))); + Thread.sleep(100); //It's here to make sure that authentication will start before the other threads. + List> requests = times( + 2, + () -> new ExternalAuthenticator(redirectHandler, onPoll(TokenPollResult::pending), KnownToken.memoryCached(), Duration.ofMillis(1)) + .authenticate(null, getUnauthorizedResponse("Bearer x_token_server=\"http://token.uri\", x_redirect_server=\"http://redirect.uri\""))) + .map(executor::submit) + .collect(toImmutableList()); + + Thread.sleep(100); + interruptableThreadPool.shutdownNow(); + + ConcurrentRequestAssertion assertion = new ConcurrentRequestAssertion(ImmutableList.>builder() + .addAll(requests) + .add(interruptedAuthentication) + .build()); + assertion.requests().containsExactly(null, null); + assertion.firstException().hasRootCauseInstanceOf(InterruptedException.class); + + assertThat(redirectHandler.getRedirectionCount()).isEqualTo(1); + } + + private static Stream> times(int times, Callable request) + { + return Stream.generate(() -> request) + .limit(times); + } + + private static Optional buildAuthentication(String challengeHeader) + { + return toAuthentication(getUnauthorizedResponse(challengeHeader)); + } + + private static Response getUnauthorizedResponse(String challengeHeader) + { + return getUnauthorizedResponse(challengeHeader, + new Request.Builder() + .url(HttpUrl.get("http://example.com")) + .build()); + } + + private static Response getUnauthorizedResponse(String challengeHeader, Request request) + { + return new Response.Builder() + .request(request) + .protocol(Protocol.HTTP_1_1) + .code(HTTP_UNAUTHORIZED) + .message("Unauthorized") + .header(WWW_AUTHENTICATE, challengeHeader) + .build(); + } + + static class ConcurrentRequestAssertion + { + private final List exceptions = new ArrayList<>(); + private final List requests = new ArrayList<>(); + + public ConcurrentRequestAssertion(List> requests) + { + for (Future request : requests) { + try { + this.requests.add(request.get()); + } + catch (InterruptedException e) { + Thread.currentThread().interrupt(); + throw new RuntimeException(e); + } + catch (CancellationException ex) { + exceptions.add(ex); + } + catch (ExecutionException ex) { + checkState(ex.getCause() != null, "Missing cause on ExecutionException " + ex.getMessage()); + + exceptions.add(ex.getCause()); + } + } + } + + ThrowableAssert firstException() + { + return exceptions.stream() + .findFirst() + .map(ThrowableAssert::new) + .orElseGet(() -> new ThrowableAssert(() -> null)); + } + + void assertThatNoExceptionsHasBeenThrown() + { + if (!exceptions.isEmpty()) { + Throwable firstException = exceptions.get(0); + AssertionError assertionError = new AssertionError("Expected no exceptions, but some exceptions has been thrown", firstException); + for (int i = 1; i < exceptions.size(); i++) { + assertionError.addSuppressed(exceptions.get(i)); + } + throw assertionError; + } + } + + ListAssert requests() + { + return assertThat(requests); + } + } +} diff --git a/presto-client/src/test/java/com/facebook/presto/client/auth/external/TestHttpTokenPoller.java b/presto-client/src/test/java/com/facebook/presto/client/auth/external/TestHttpTokenPoller.java new file mode 100644 index 0000000000000..69f110d4ec917 --- /dev/null +++ b/presto-client/src/test/java/com/facebook/presto/client/auth/external/TestHttpTokenPoller.java @@ -0,0 +1,224 @@ +/* + * Licensed 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 com.facebook.presto.client.auth.external; + +import okhttp3.HttpUrl; +import okhttp3.OkHttpClient; +import okhttp3.mockwebserver.MockResponse; +import okhttp3.mockwebserver.MockWebServer; +import okhttp3.mockwebserver.RecordedRequest; +import org.testng.annotations.AfterMethod; +import org.testng.annotations.BeforeMethod; +import org.testng.annotations.Test; + +import java.io.IOException; +import java.io.UncheckedIOException; +import java.net.URI; +import java.time.Duration; + +import static com.google.common.net.HttpHeaders.CONTENT_TYPE; +import static com.google.common.net.MediaType.JSON_UTF_8; +import static java.lang.String.format; +import static java.net.HttpURLConnection.HTTP_GONE; +import static java.net.HttpURLConnection.HTTP_OK; +import static java.net.HttpURLConnection.HTTP_UNAVAILABLE; +import static java.net.URI.create; +import static java.util.concurrent.TimeUnit.MILLISECONDS; +import static org.assertj.core.api.Assertions.assertThat; +import static org.assertj.core.api.Assertions.assertThatThrownBy; + +@Test(singleThreaded = true) +public class TestHttpTokenPoller +{ + private static final String TOKEN_PATH = "/v1/authentications/sso/test/token"; + private static final Duration ONE_SECOND = Duration.ofSeconds(1); + + private TokenPoller tokenPoller; + private MockWebServer server; + + @BeforeMethod(alwaysRun = true) + public void setup() + throws Exception + { + server = new MockWebServer(); + server.start(); + + tokenPoller = new HttpTokenPoller(new OkHttpClient.Builder() + .callTimeout(Duration.ofMillis(500)) + .build()); + } + + @AfterMethod(alwaysRun = true) + public void teardown() + throws IOException + { + server.close(); + server = null; + } + + @Test + public void testTokenReady() + { + server.enqueue(statusAndBody(HTTP_OK, jsonPair("token", "token"))); + + TokenPollResult result = tokenPoller.pollForToken(tokenUri(), ONE_SECOND); + + assertThat(result.getToken().token()).isEqualTo("token"); + assertThat(server.getRequestCount()).isEqualTo(1); + } + + @Test + public void testTokenNotReady() + { + server.enqueue(statusAndBody(HTTP_OK, jsonPair("nextUri", tokenUri()))); + + TokenPollResult result = tokenPoller.pollForToken(tokenUri(), ONE_SECOND); + + assertThat(result.isPending()).isTrue(); + assertThat(server.getRequestCount()).isEqualTo(1); + } + + @Test + public void testErrorResponse() + { + server.enqueue(statusAndBody(HTTP_OK, jsonPair("error", "test failure"))); + + TokenPollResult result = tokenPoller.pollForToken(tokenUri(), ONE_SECOND); + + assertThat(result.isFailed()).isTrue(); + assertThat(result.getError()).contains("test failure"); + } + + @Test + public void testBadHttpStatus() + { + server.enqueue(new MockResponse().setResponseCode(HTTP_GONE)); + + TokenPollResult result = tokenPoller.pollForToken(tokenUri(), ONE_SECOND); + + assertThat(result.isFailed()).isTrue(); + assertThat(result.getError()) + .matches("Request to http://.* failed: JsonResponse\\{statusCode=410, .*"); + } + + @Test + public void testInvalidJsonBody() + { + server.enqueue(statusAndBody(HTTP_OK, jsonPair("foo", "bar"))); + + TokenPollResult result = tokenPoller.pollForToken(tokenUri(), ONE_SECOND); + + assertThat(result.isFailed()).isTrue(); + assertThat(result.getError()) + .isEqualTo("Failed to poll for token. No fields set in response."); + } + + @Test + public void testInvalidNextUri() + { + server.enqueue(statusAndBody(HTTP_OK, jsonPair("nextUri", ":::"))); + + TokenPollResult result = tokenPoller.pollForToken(tokenUri(), ONE_SECOND); + + assertThat(result.isFailed()).isTrue(); + assertThat(result.getError()) + .matches("Request to http://.* failed: JsonResponse\\{statusCode=200, .*, hasValue=false} .*"); + } + + @Test + public void testHttpStatus503() + { + for (int i = 1; i <= 100; i++) { + server.enqueue(statusAndBody(HTTP_UNAVAILABLE, "Server failure #" + i)); + } + + assertThatThrownBy(() -> tokenPoller.pollForToken(tokenUri(), ONE_SECOND)) + .isInstanceOf(UncheckedIOException.class) + .hasRootCauseExactlyInstanceOf(IOException.class); + + assertThat(server.getRequestCount()).isGreaterThan(1); + } + + @Test + public void testHttpTimeout() + { + // force request to timeout by not enqueuing response + + assertThatThrownBy(() -> tokenPoller.pollForToken(tokenUri(), ONE_SECOND)) + .isInstanceOf(UncheckedIOException.class) + .hasMessageEndingWith(": timeout"); + } + + @Test + public void testTokenReceived() + throws InterruptedException + { + server.enqueue(status(HTTP_OK)); + + tokenPoller.tokenReceived(tokenUri()); + + RecordedRequest request = server.takeRequest(1, MILLISECONDS); + assertThat(request.getMethod()).isEqualTo("DELETE"); + assertThat(request.getRequestUrl()).isEqualTo(HttpUrl.get(tokenUri())); + } + + @Test + public void testTokenReceivedRetriesUntilNotErrorReturned() + { + server.enqueue(status(HTTP_UNAVAILABLE)); + server.enqueue(status(HTTP_UNAVAILABLE)); + server.enqueue(status(HTTP_UNAVAILABLE)); + server.enqueue(status(202)); + + tokenPoller.tokenReceived(tokenUri()); + + assertThat(server.getRequestCount()).isEqualTo(4); + } + + @Test + public void testTokenReceivedDoesNotRetriesIndefinitely() + { + for (int i = 1; i <= 100; i++) { + server.enqueue(status(HTTP_UNAVAILABLE)); + } + + tokenPoller.tokenReceived(tokenUri()); + + assertThat(server.getRequestCount()).isLessThan(100); + } + + private URI tokenUri() + { + return create("http://" + server.getHostName() + ":" + server.getPort() + TOKEN_PATH); + } + + private static String jsonPair(String key, Object value) + { + return format("{\"%s\": \"%s\"}", key, value); + } + + private static MockResponse statusAndBody(int status, String body) + { + return new MockResponse() + .setResponseCode(status) + .addHeader(CONTENT_TYPE, JSON_UTF_8) + .setBody(body); + } + + private static MockResponse status(int status) + { + return new MockResponse() + .setResponseCode(status); + } +} diff --git a/presto-common/src/main/java/com/facebook/presto/common/SourceColumn.java b/presto-common/src/main/java/com/facebook/presto/common/SourceColumn.java new file mode 100644 index 0000000000000..9125e3e120c73 --- /dev/null +++ b/presto-common/src/main/java/com/facebook/presto/common/SourceColumn.java @@ -0,0 +1,77 @@ +/* + * Licensed 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 com.facebook.presto.common; + +import com.fasterxml.jackson.annotation.JsonCreator; +import com.fasterxml.jackson.annotation.JsonProperty; + +import java.util.Objects; + +import static java.util.Objects.requireNonNull; + +public class SourceColumn +{ + private final QualifiedObjectName tableName; + private final String columnName; + + @JsonCreator + public SourceColumn( + @JsonProperty("tableName") QualifiedObjectName tableName, + @JsonProperty("columnName") String columnName) + { + this.tableName = requireNonNull(tableName, "tableName is null"); + this.columnName = requireNonNull(columnName, "columnName is null"); + } + + @JsonProperty + public QualifiedObjectName getTableName() + { + return tableName; + } + + @JsonProperty + public String getColumnName() + { + return columnName; + } + + @Override + public int hashCode() + { + return Objects.hash(tableName, columnName); + } + + @Override + public boolean equals(Object obj) + { + if (obj == this) { + return true; + } + if ((obj == null) || (getClass() != obj.getClass())) { + return false; + } + SourceColumn entry = (SourceColumn) obj; + return Objects.equals(tableName, entry.tableName) && + Objects.equals(columnName, entry.columnName); + } + + @Override + public String toString() + { + return "SourceColumn{" + + "tableName=" + tableName + + ", columnName='" + columnName + '\'' + + '}'; + } +} diff --git a/presto-common/src/main/java/com/facebook/presto/common/type/TypeManager.java b/presto-common/src/main/java/com/facebook/presto/common/type/TypeManager.java index 277ef86ffb881..2dc2aa18cfaa2 100644 --- a/presto-common/src/main/java/com/facebook/presto/common/type/TypeManager.java +++ b/presto-common/src/main/java/com/facebook/presto/common/type/TypeManager.java @@ -44,4 +44,9 @@ default Collection getParametricTypes() { throw new UnsupportedOperationException(); } + + /** + * Checks for the existence of this type. + */ + boolean hasType(TypeSignature signature); } diff --git a/presto-common/src/test/java/com/facebook/presto/common/type/TestingTypeManager.java b/presto-common/src/test/java/com/facebook/presto/common/type/TestingTypeManager.java index f2e6f847a48a5..e94a0a9465cf0 100644 --- a/presto-common/src/test/java/com/facebook/presto/common/type/TestingTypeManager.java +++ b/presto-common/src/test/java/com/facebook/presto/common/type/TestingTypeManager.java @@ -59,4 +59,10 @@ public List getTypes() { return ImmutableList.of(BOOLEAN, INTEGER, BIGINT, DOUBLE, VARCHAR, VARBINARY, TIMESTAMP, DATE, ID, HYPER_LOG_LOG); } + + @Override + public boolean hasType(TypeSignature signature) + { + return getType(signature) != null; + } } diff --git a/presto-db-session-property-manager/pom.xml b/presto-db-session-property-manager/pom.xml new file mode 100644 index 0000000000000..d0ef4ca921014 --- /dev/null +++ b/presto-db-session-property-manager/pom.xml @@ -0,0 +1,175 @@ + + + 4.0.0 + + + com.facebook.presto + presto-root + 0.295-SNAPSHOT + + + presto-db-session-property-manager + presto-db-session-property-manager + Presto - DB Session Property Manager + presto-plugin + + + ${project.parent.basedir} + + + + + com.facebook.presto + presto-session-property-managers-common + + + + com.facebook.airlift + bootstrap + + + + com.facebook.airlift + log + + + + com.facebook.airlift + json + + + + com.facebook.airlift + configuration + + + + com.google.guava + guava + + + + com.google.inject + guice + + + + javax.annotation + javax.annotation-api + + + + javax.inject + javax.inject + + + + com.facebook.airlift + concurrent + + + + jakarta.validation + jakarta.validation-api + + + + org.jdbi + jdbi3-core + + + + org.jdbi + jdbi3-sqlobject + + + + com.mysql + mysql-connector-j + true + runtime + + + + org.mariadb.jdbc + mariadb-java-client + true + runtime + + + + + com.facebook.presto + presto-spi + provided + + + + com.facebook.presto + presto-common + provided + + + + com.facebook.airlift + units + provided + + + + io.airlift + slice + provided + + + + com.fasterxml.jackson.core + jackson-annotations + provided + + + + org.openjdk.jol + jol-core + provided + + + + + com.facebook.presto + presto-testng-services + test + + + + org.testng + testng + test + + + + com.facebook.airlift + testing + test + + + + com.facebook.presto + testing-mysql-server-8 + test + + + + com.facebook.presto + testing-mysql-server-base + test + + + + com.facebook.presto + presto-session-property-managers-common + test-jar + test + + + diff --git a/presto-session-property-managers/src/main/java/com/facebook/presto/session/FileSessionPropertyManagerPlugin.java b/presto-db-session-property-manager/src/main/java/com/facebook/presto/session/db/DbSessionPropertyConfigurationManagerPlugin.java similarity index 82% rename from presto-session-property-managers/src/main/java/com/facebook/presto/session/FileSessionPropertyManagerPlugin.java rename to presto-db-session-property-manager/src/main/java/com/facebook/presto/session/db/DbSessionPropertyConfigurationManagerPlugin.java index bde86e9400f31..bbbc653abd2c9 100644 --- a/presto-session-property-managers/src/main/java/com/facebook/presto/session/FileSessionPropertyManagerPlugin.java +++ b/presto-db-session-property-manager/src/main/java/com/facebook/presto/session/db/DbSessionPropertyConfigurationManagerPlugin.java @@ -11,18 +11,19 @@ * See the License for the specific language governing permissions and * limitations under the License. */ -package com.facebook.presto.session; +package com.facebook.presto.session.db; import com.facebook.presto.spi.Plugin; import com.facebook.presto.spi.session.SessionPropertyConfigurationManagerFactory; import com.google.common.collect.ImmutableList; -public class FileSessionPropertyManagerPlugin +public class DbSessionPropertyConfigurationManagerPlugin implements Plugin { @Override public Iterable getSessionPropertyConfigurationManagerFactories() { - return ImmutableList.of(new FileSessionPropertyManagerFactory()); + return ImmutableList.of( + new DbSessionPropertyManagerFactory()); } } diff --git a/presto-db-session-property-manager/src/main/java/com/facebook/presto/session/db/DbSessionPropertyManager.java b/presto-db-session-property-manager/src/main/java/com/facebook/presto/session/db/DbSessionPropertyManager.java new file mode 100644 index 0000000000000..528d51c395c6e --- /dev/null +++ b/presto-db-session-property-manager/src/main/java/com/facebook/presto/session/db/DbSessionPropertyManager.java @@ -0,0 +1,47 @@ +/* + * Licensed 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 com.facebook.presto.session.db; + +import com.facebook.presto.session.AbstractSessionPropertyManager; +import com.facebook.presto.session.SessionMatchSpec; +import com.facebook.presto.spi.session.SessionConfigurationContext; +import com.facebook.presto.spi.session.SessionPropertyConfigurationManager; + +import javax.inject.Inject; + +import java.util.List; + +import static java.util.Objects.requireNonNull; + +/** + * A {@link SessionPropertyConfigurationManager} implementation that connects to a database for fetching information + * about session property overrides given {@link SessionConfigurationContext}. + */ +public class DbSessionPropertyManager + extends AbstractSessionPropertyManager +{ + private final DbSpecsProvider specsProvider; + + @Inject + public DbSessionPropertyManager(DbSpecsProvider specsProvider) + { + this.specsProvider = requireNonNull(specsProvider, "specsProvider is null"); + } + + @Override + protected List getSessionMatchSpecs() + { + return this.specsProvider.get(); + } +} diff --git a/presto-db-session-property-manager/src/main/java/com/facebook/presto/session/db/DbSessionPropertyManagerConfig.java b/presto-db-session-property-manager/src/main/java/com/facebook/presto/session/db/DbSessionPropertyManagerConfig.java new file mode 100644 index 0000000000000..4a94ac0d9e61a --- /dev/null +++ b/presto-db-session-property-manager/src/main/java/com/facebook/presto/session/db/DbSessionPropertyManagerConfig.java @@ -0,0 +1,54 @@ +/* + * Licensed 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 com.facebook.presto.session.db; + +import com.facebook.airlift.configuration.Config; +import com.facebook.airlift.units.Duration; +import com.facebook.airlift.units.MinDuration; +import jakarta.validation.constraints.NotNull; + +import static java.util.concurrent.TimeUnit.SECONDS; + +public class DbSessionPropertyManagerConfig +{ + private String configDbUrl; + private Duration specsRefreshPeriod = new Duration(10, SECONDS); + + @NotNull + public String getConfigDbUrl() + { + return configDbUrl; + } + + @Config("session-property-manager.db.url") + public DbSessionPropertyManagerConfig setConfigDbUrl(String configDbUrl) + { + this.configDbUrl = configDbUrl; + return this; + } + + @NotNull + @MinDuration("1ms") + public Duration getSpecsRefreshPeriod() + { + return specsRefreshPeriod; + } + + @Config("session-property-manager.db.refresh-period") + public DbSessionPropertyManagerConfig setSpecsRefreshPeriod(Duration specsRefreshPeriod) + { + this.specsRefreshPeriod = specsRefreshPeriod; + return this; + } +} diff --git a/presto-db-session-property-manager/src/main/java/com/facebook/presto/session/db/DbSessionPropertyManagerFactory.java b/presto-db-session-property-manager/src/main/java/com/facebook/presto/session/db/DbSessionPropertyManagerFactory.java new file mode 100644 index 0000000000000..7a6eba7821ef0 --- /dev/null +++ b/presto-db-session-property-manager/src/main/java/com/facebook/presto/session/db/DbSessionPropertyManagerFactory.java @@ -0,0 +1,53 @@ +/* + * Licensed 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 com.facebook.presto.session.db; + +import com.facebook.airlift.bootstrap.Bootstrap; +import com.facebook.airlift.json.JsonModule; +import com.facebook.presto.spi.resourceGroups.SessionPropertyConfigurationManagerContext; +import com.facebook.presto.spi.session.SessionPropertyConfigurationManager; +import com.facebook.presto.spi.session.SessionPropertyConfigurationManagerFactory; +import com.google.inject.Injector; + +import java.util.Map; + +import static com.google.common.base.Throwables.throwIfUnchecked; + +public class DbSessionPropertyManagerFactory + implements SessionPropertyConfigurationManagerFactory +{ + @Override + public String getName() + { + return "db"; + } + + @Override + public SessionPropertyConfigurationManager create(Map config, SessionPropertyConfigurationManagerContext context) + { + try { + Bootstrap app = new Bootstrap(new JsonModule(), new DbSessionPropertyManagerModule()); + Injector injector = app + .doNotInitializeLogging() + .setRequiredConfigurationProperties(config) + .initialize(); + + return injector.getInstance(DbSessionPropertyManager.class); + } + catch (Exception e) { + throwIfUnchecked(e); + throw new RuntimeException(e); + } + } +} diff --git a/presto-db-session-property-manager/src/main/java/com/facebook/presto/session/db/DbSessionPropertyManagerModule.java b/presto-db-session-property-manager/src/main/java/com/facebook/presto/session/db/DbSessionPropertyManagerModule.java new file mode 100644 index 0000000000000..48f480b7603aa --- /dev/null +++ b/presto-db-session-property-manager/src/main/java/com/facebook/presto/session/db/DbSessionPropertyManagerModule.java @@ -0,0 +1,33 @@ +/* + * Licensed 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 com.facebook.presto.session.db; + +import com.google.inject.Binder; +import com.google.inject.Module; +import com.google.inject.Scopes; + +import static com.facebook.airlift.configuration.ConfigBinder.configBinder; + +public class DbSessionPropertyManagerModule + implements Module +{ + @Override + public void configure(Binder binder) + { + configBinder(binder).bindConfig(DbSessionPropertyManagerConfig.class); + binder.bind(DbSessionPropertyManager.class).in(Scopes.SINGLETON); + binder.bind(SessionPropertiesDao.class).toProvider(SessionPropertiesDaoProvider.class).in(Scopes.SINGLETON); + binder.bind(DbSpecsProvider.class).to(RefreshingDbSpecsProvider.class).in(Scopes.SINGLETON); + } +} diff --git a/presto-db-session-property-manager/src/main/java/com/facebook/presto/session/db/DbSpecsProvider.java b/presto-db-session-property-manager/src/main/java/com/facebook/presto/session/db/DbSpecsProvider.java new file mode 100644 index 0000000000000..278b52daadb67 --- /dev/null +++ b/presto-db-session-property-manager/src/main/java/com/facebook/presto/session/db/DbSpecsProvider.java @@ -0,0 +1,29 @@ +/* + * Licensed 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 com.facebook.presto.session.db; + +import com.facebook.presto.session.SessionMatchSpec; + +import java.util.List; +import java.util.function.Supplier; + +/** + * This interface was created to separate the scheduling logic for {@link SessionMatchSpec} loading. This also helps + * us test the core logic of {@link DbSessionPropertyManager} in a modular fashion by letting us use a test + * implementation of this interface. + */ +public interface DbSpecsProvider + extends Supplier> +{ +} diff --git a/presto-db-session-property-manager/src/main/java/com/facebook/presto/session/db/RefreshingDbSpecsProvider.java b/presto-db-session-property-manager/src/main/java/com/facebook/presto/session/db/RefreshingDbSpecsProvider.java new file mode 100644 index 0000000000000..2694fd5e54e9e --- /dev/null +++ b/presto-db-session-property-manager/src/main/java/com/facebook/presto/session/db/RefreshingDbSpecsProvider.java @@ -0,0 +1,94 @@ +/* + * Licensed 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 com.facebook.presto.session.db; + +import com.facebook.airlift.log.Logger; +import com.facebook.presto.session.SessionMatchSpec; +import com.google.common.annotations.VisibleForTesting; +import com.google.common.collect.ImmutableList; + +import javax.annotation.PostConstruct; +import javax.annotation.PreDestroy; +import javax.inject.Inject; + +import java.util.List; +import java.util.concurrent.ScheduledExecutorService; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicBoolean; +import java.util.concurrent.atomic.AtomicReference; + +import static com.facebook.airlift.concurrent.Threads.daemonThreadsNamed; +import static java.util.Objects.requireNonNull; +import static java.util.concurrent.Executors.newSingleThreadScheduledExecutor; + +/** + * Periodically schedules the loading of specs from the database during initialization. Returns the most recent successfully + * loaded specs on every get() invocation. + */ +public class RefreshingDbSpecsProvider + implements DbSpecsProvider +{ + private static final Logger log = Logger.get(RefreshingDbSpecsProvider.class); + + private final AtomicReference> sessionMatchSpecs = new AtomicReference<>(ImmutableList.of()); + private final SessionPropertiesDao dao; + + private final ScheduledExecutorService executor = newSingleThreadScheduledExecutor(daemonThreadsNamed("RefreshingDbSpecsProvider")); + private final AtomicBoolean started = new AtomicBoolean(); + private final long refreshPeriodMillis; + + @Inject + public RefreshingDbSpecsProvider(DbSessionPropertyManagerConfig config, SessionPropertiesDao dao) + { + requireNonNull(config, "config is null"); + this.dao = requireNonNull(dao, "dao is null"); + this.refreshPeriodMillis = config.getSpecsRefreshPeriod().toMillis(); + dao.createSessionSpecsTable(); + dao.createSessionClientTagsTable(); + dao.createSessionPropertiesTable(); + } + + @PostConstruct + public void initialize() + { + if (!started.getAndSet(true)) { + executor.scheduleWithFixedDelay(this::refresh, 0, refreshPeriodMillis, TimeUnit.MILLISECONDS); + } + } + + @VisibleForTesting + void refresh() + { + try { + sessionMatchSpecs.set(ImmutableList.copyOf(dao.getSessionMatchSpecs())); + } + catch (Throwable e) { + // Catch all exceptions here since throwing an exception from executor#scheduleWithFixedDelay method + // suppresses all future scheduled invocations + log.error(e, "Error loading configuration from database"); + } + } + + @PreDestroy + public void destroy() + { + executor.shutdownNow(); + } + + @Override + public List get() + { + return sessionMatchSpecs.get(); + } +} diff --git a/presto-db-session-property-manager/src/main/java/com/facebook/presto/session/db/SessionPropertiesDao.java b/presto-db-session-property-manager/src/main/java/com/facebook/presto/session/db/SessionPropertiesDao.java new file mode 100644 index 0000000000000..71a93f0838c4a --- /dev/null +++ b/presto-db-session-property-manager/src/main/java/com/facebook/presto/session/db/SessionPropertiesDao.java @@ -0,0 +1,130 @@ +/* + * Licensed 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 com.facebook.presto.session.db; + +import com.facebook.presto.session.SessionMatchSpec; +import com.google.common.annotations.VisibleForTesting; +import org.jdbi.v3.sqlobject.customizer.Bind; +import org.jdbi.v3.sqlobject.statement.SqlQuery; +import org.jdbi.v3.sqlobject.statement.SqlUpdate; +import org.jdbi.v3.sqlobject.statement.UseRowMapper; + +import java.util.List; + +/** + * Dao should guarantee that the list of SessionMatchSpecs is returned in increasing order of priority. i.e. if two + * rows in the ResultSet specify different values for the same property, the row coming in later will override the + * value set by the row coming in earlier. + */ +public interface SessionPropertiesDao +{ + String SESSION_SPECS_TABLE = "session_specs"; + String CLIENT_TAGS_TABLE = "session_client_tags"; + String PROPERTIES_TABLE = "session_property_values"; + public static String EMPTY_CATALOG = "__NULL__"; + + @SqlUpdate("CREATE TABLE IF NOT EXISTS " + SESSION_SPECS_TABLE + "(\n" + + "spec_id BIGINT NOT NULL AUTO_INCREMENT,\n" + + "user_regex VARCHAR(512),\n" + + "source_regex VARCHAR(512),\n" + + "query_type VARCHAR(512),\n" + + "group_regex VARCHAR(512),\n" + + "client_info_regex VARCHAR(512),\n" + + "override_session_properties TINYINT(1),\n" + + "priority INT NOT NULL,\n" + + "PRIMARY KEY (spec_id)\n" + + ")") + void createSessionSpecsTable(); + + @SqlUpdate("CREATE TABLE IF NOT EXISTS " + CLIENT_TAGS_TABLE + "(\n" + + "tag_spec_id BIGINT NOT NULL,\n" + + "client_tag VARCHAR(512) NOT NULL,\n" + + "PRIMARY KEY (tag_spec_id, client_tag),\n" + + "FOREIGN KEY (tag_spec_id) REFERENCES session_specs (spec_id)\n" + + ")") + void createSessionClientTagsTable(); + + @SqlUpdate("CREATE TABLE IF NOT EXISTS " + PROPERTIES_TABLE + "(\n" + + "property_spec_id BIGINT NOT NULL,\n" + + "session_property_name VARCHAR(512),\n" + + "session_property_value VARCHAR(512),\n" + + "catalog VARCHAR(512),\n" + + "PRIMARY KEY (property_spec_id, session_property_name),\n" + + "FOREIGN KEY (property_spec_id) REFERENCES session_specs (spec_id)\n" + + ")") + void createSessionPropertiesTable(); + + @SqlUpdate("DROP TABLE IF EXISTS " + SESSION_SPECS_TABLE) + void dropSessionSpecsTable(); + + @SqlUpdate("DROP TABLE IF EXISTS " + CLIENT_TAGS_TABLE) + void dropSessionClientTagsTable(); + + @SqlUpdate("DROP TABLE IF EXISTS " + PROPERTIES_TABLE) + void dropSessionPropertiesTable(); + + @SqlQuery("SELECT " + + "S.spec_id,\n" + + "S.user_regex,\n" + + "S.source_regex,\n" + + "S.query_type,\n" + + "S.group_regex,\n" + + "S.client_info_regex,\n" + + "S.override_session_properties,\n" + + "S.client_tags,\n" + + "GROUP_CONCAT(P.session_property_name ORDER BY P.session_property_name) session_property_names,\n" + + "GROUP_CONCAT(P.session_property_value ORDER BY P.session_property_name) session_property_values,\n" + + "GROUP_CONCAT(COALESCE(P.catalog, '" + EMPTY_CATALOG + "') ORDER BY P.session_property_name) session_property_catalogs\n" + + "FROM\n" + + "(SELECT\n" + + "A.spec_id, A.user_regex, A.source_regex, A.query_type, A.group_regex, A.client_info_regex, A.override_session_properties, A.priority,\n" + + "GROUP_CONCAT(DISTINCT B.client_tag) client_tags\n" + + "FROM " + SESSION_SPECS_TABLE + " A\n" + + "LEFT JOIN " + CLIENT_TAGS_TABLE + " B\n" + + "ON A.spec_id = B.tag_spec_id\n" + + "GROUP BY A.spec_id, A.user_regex, A.source_regex, A.query_type, A.group_regex, A.client_info_regex, A.override_session_properties, A.priority)\n" + + " S JOIN\n" + + PROPERTIES_TABLE + " P\n" + + "ON S.spec_id = P.property_spec_id\n" + + "GROUP BY S.spec_id, S.user_regex, S.source_regex, S.query_type, S.group_regex, S.client_info_regex, S.override_session_properties, S.priority, S.client_tags\n" + + "ORDER BY S.priority asc") + @UseRowMapper(SessionMatchSpec.Mapper.class) + List getSessionMatchSpecs(); + + @VisibleForTesting + @SqlUpdate("INSERT INTO " + SESSION_SPECS_TABLE + " (spec_id, user_regex, source_regex, query_type, group_regex, client_info_regex, override_session_properties, priority)\n" + + "VALUES (:spec_id, :user_regex, :source_regex, :query_type, :group_regex, :client_info_regex, :override_session_properties, :priority)") + void insertSpecRow( + @Bind("spec_id") long specId, + @Bind("user_regex") String userRegex, + @Bind("source_regex") String sourceRegex, + @Bind("query_type") String queryType, + @Bind("group_regex") String groupRegex, + @Bind("client_info_regex") String clientInfoRegex, + @Bind("override_session_properties") Integer overrideSessionProperties, + @Bind("priority") int priority); + + @VisibleForTesting + @SqlUpdate("INSERT INTO " + CLIENT_TAGS_TABLE + " (tag_spec_id, client_tag) VALUES (:spec_id, :client_tag)") + void insertClientTag(@Bind("spec_id") long specId, @Bind("client_tag") String clientTag); + + @VisibleForTesting + @SqlUpdate("INSERT INTO " + PROPERTIES_TABLE + " (property_spec_id, session_property_name, session_property_value, catalog)\n" + + "VALUES (:property_spec_id, :session_property_name, :session_property_value, :catalog)") + void insertSessionProperty( + @Bind("property_spec_id") long propertySpecId, + @Bind("session_property_name") String sessionPropertyName, + @Bind("session_property_value") String sessionPropertyValue, + @Bind("catalog") String catalog); +} diff --git a/presto-db-session-property-manager/src/main/java/com/facebook/presto/session/db/SessionPropertiesDaoProvider.java b/presto-db-session-property-manager/src/main/java/com/facebook/presto/session/db/SessionPropertiesDaoProvider.java new file mode 100644 index 0000000000000..06617a1885396 --- /dev/null +++ b/presto-db-session-property-manager/src/main/java/com/facebook/presto/session/db/SessionPropertiesDaoProvider.java @@ -0,0 +1,46 @@ +/* + * Licensed 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 com.facebook.presto.session.db; + +import org.jdbi.v3.core.Jdbi; +import org.jdbi.v3.sqlobject.SqlObjectPlugin; + +import javax.inject.Inject; +import javax.inject.Provider; + +import java.sql.DriverManager; + +import static java.util.Objects.requireNonNull; + +public class SessionPropertiesDaoProvider + implements Provider +{ + private final SessionPropertiesDao dao; + + @Inject + public SessionPropertiesDaoProvider(DbSessionPropertyManagerConfig config) + { + requireNonNull(config, "config is null"); + requireNonNull(config.getConfigDbUrl(), "db url is null"); + this.dao = Jdbi.create(() -> DriverManager.getConnection(config.getConfigDbUrl())) + .installPlugin(new SqlObjectPlugin()) + .onDemand(SessionPropertiesDao.class); + } + + @Override + public SessionPropertiesDao get() + { + return dao; + } +} diff --git a/presto-db-session-property-manager/src/test/java/com/facebook/presto/session/db/TestDbSessionPropertyManager.java b/presto-db-session-property-manager/src/test/java/com/facebook/presto/session/db/TestDbSessionPropertyManager.java new file mode 100644 index 0000000000000..865e8f840edf4 --- /dev/null +++ b/presto-db-session-property-manager/src/test/java/com/facebook/presto/session/db/TestDbSessionPropertyManager.java @@ -0,0 +1,112 @@ +/* + * Licensed 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 com.facebook.presto.session.db; + +import com.facebook.presto.session.AbstractTestSessionPropertyManager; +import com.facebook.presto.session.SessionMatchSpec; +import com.facebook.presto.spi.session.SessionPropertyConfigurationManager; +import com.facebook.presto.testing.mysql.MySqlOptions; +import com.facebook.presto.testing.mysql.TestingMySqlServer; +import com.google.common.collect.ImmutableList; +import com.google.common.collect.ImmutableMap; +import org.testng.annotations.AfterClass; + +import java.util.Map; +import java.util.regex.Pattern; + +import static org.testng.Assert.assertEquals; + +public abstract class TestDbSessionPropertyManager + extends AbstractTestSessionPropertyManager +{ + private static final MySqlOptions MY_SQL_OPTIONS = MySqlOptions.builder() + .build(); + + private final String driver; + + private final TestingMySqlServer mysqlServer; + + public TestDbSessionPropertyManager(String driver) + throws Exception + { + this.driver = driver; + this.mysqlServer = new TestingMySqlServer("testuser", "testpass", ImmutableList.of(), MY_SQL_OPTIONS); + } + + @AfterClass(alwaysRun = true) + public void tearDown() + throws Exception + { + mysqlServer.close(); + } + + @Override + protected void assertProperties(Map defaultProperties, SessionMatchSpec... spec) + { + assertProperties(defaultProperties, ImmutableMap.of(), ImmutableMap.of(), spec); + } + + @Override + protected void assertProperties(Map defaultProperties, Map overrideProperties, SessionMatchSpec... specs) + { + assertProperties(defaultProperties, overrideProperties, ImmutableMap.of(), specs); + } + + @Override + protected void assertProperties(Map defaultProperties, Map overrideProperties, Map> catalogProperties, SessionMatchSpec... specs) + { + DbSessionPropertyManagerConfig config = new DbSessionPropertyManagerConfig() + .setConfigDbUrl(overrideJdbcUrl(mysqlServer.getJdbcUrl("session") + "&createDatabaseIfNotExist=true")); + + SessionPropertiesDaoProvider sessionPropertiesDaoProvider = new SessionPropertiesDaoProvider(config); + SessionPropertiesDao dao = sessionPropertiesDaoProvider.get(); + RefreshingDbSpecsProvider dbSpecsProvider = new RefreshingDbSpecsProvider(config, sessionPropertiesDaoProvider.get()); + SessionPropertyConfigurationManager manager = new DbSessionPropertyManager(dbSpecsProvider); + int id = 1; + try { + for (SessionMatchSpec spec : specs) { + int finalId = id; + dao.insertSpecRow( + finalId, + spec.getUserRegex().map(Pattern::pattern).orElse(null), + spec.getSourceRegex().map(Pattern::pattern).orElse(null), + spec.getQueryType().orElse(null), + spec.getResourceGroupRegex().map(Pattern::pattern).orElse(null), + spec.getClientInfoRegex().map(Pattern::pattern).orElse(null), + spec.getOverrideSessionProperties().map(val -> val ? 1 : 0).orElse(null), + finalId); + spec.getClientTags().forEach(tag -> dao.insertClientTag(finalId, tag)); + spec.getSessionProperties().forEach((key, value) -> dao.insertSessionProperty(finalId, key, value, null)); + spec.getCatalogSessionProperties().forEach((catalog, property) -> property.forEach((key, value) -> dao.insertSessionProperty(finalId, key, value, catalog))); + id++; + } + dbSpecsProvider.refresh(); + SessionPropertyConfigurationManager.SystemSessionPropertyConfiguration propertyConfiguration = manager.getSystemSessionProperties(CONTEXT); + assertEquals(propertyConfiguration.systemPropertyDefaults, defaultProperties); + assertEquals(propertyConfiguration.systemPropertyOverrides, overrideProperties); + assertEquals(manager.getCatalogSessionProperties(CONTEXT), catalogProperties); + } + finally { + dao.dropSessionPropertiesTable(); + dao.dropSessionClientTagsTable(); + dao.dropSessionSpecsTable(); + dbSpecsProvider.destroy(); + } + } + + protected String overrideJdbcUrl(String url) + { + return url; + } +} diff --git a/presto-db-session-property-manager/src/test/java/com/facebook/presto/session/db/TestDbSessionPropertyManagerConfig.java b/presto-db-session-property-manager/src/test/java/com/facebook/presto/session/db/TestDbSessionPropertyManagerConfig.java new file mode 100644 index 0000000000000..d976eb722390c --- /dev/null +++ b/presto-db-session-property-manager/src/test/java/com/facebook/presto/session/db/TestDbSessionPropertyManagerConfig.java @@ -0,0 +1,52 @@ +/* + * Licensed 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 com.facebook.presto.session.db; + +import com.facebook.airlift.units.Duration; +import com.google.common.collect.ImmutableMap; +import org.testng.annotations.Test; + +import java.util.Map; +import java.util.concurrent.TimeUnit; + +import static com.facebook.airlift.configuration.testing.ConfigAssertions.assertFullMapping; +import static com.facebook.airlift.configuration.testing.ConfigAssertions.assertRecordedDefaults; +import static com.facebook.airlift.configuration.testing.ConfigAssertions.recordDefaults; +import static java.util.concurrent.TimeUnit.SECONDS; + +public class TestDbSessionPropertyManagerConfig +{ + @Test + public void testDefaults() + { + assertRecordedDefaults(recordDefaults(DbSessionPropertyManagerConfig.class) + .setConfigDbUrl(null) + .setSpecsRefreshPeriod(new Duration(10, SECONDS))); + } + + @Test + public void testExplicitPropertyMappings() + { + Map properties = new ImmutableMap.Builder() + .put("session-property-manager.db.url", "foo") + .put("session-property-manager.db.refresh-period", "50s") + .build(); + + DbSessionPropertyManagerConfig expected = new DbSessionPropertyManagerConfig() + .setConfigDbUrl("foo") + .setSpecsRefreshPeriod(new Duration(50, TimeUnit.SECONDS)); + + assertFullMapping(properties, expected); + } +} diff --git a/presto-db-session-property-manager/src/test/java/com/facebook/presto/session/db/TestDbSessionPropertyManagerMariadb.java b/presto-db-session-property-manager/src/test/java/com/facebook/presto/session/db/TestDbSessionPropertyManagerMariadb.java new file mode 100644 index 0000000000000..3e440b95d3070 --- /dev/null +++ b/presto-db-session-property-manager/src/test/java/com/facebook/presto/session/db/TestDbSessionPropertyManagerMariadb.java @@ -0,0 +1,29 @@ +/* + * Licensed 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 com.facebook.presto.session.db; + +public class TestDbSessionPropertyManagerMariadb + extends TestDbSessionPropertyManager +{ + public TestDbSessionPropertyManagerMariadb() throws Exception + { + super("mariadb"); + } + + @Override + public String overrideJdbcUrl(String url) + { + return url.replaceFirst("jdbc:mysql:", "jdbc:mariadb:"); + } +} diff --git a/presto-db-session-property-manager/src/test/java/com/facebook/presto/session/db/TestDbSessionPropertyManagerMysql.java b/presto-db-session-property-manager/src/test/java/com/facebook/presto/session/db/TestDbSessionPropertyManagerMysql.java new file mode 100644 index 0000000000000..806f5c8da9cac --- /dev/null +++ b/presto-db-session-property-manager/src/test/java/com/facebook/presto/session/db/TestDbSessionPropertyManagerMysql.java @@ -0,0 +1,23 @@ +/* + * Licensed 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 com.facebook.presto.session.db; + +public class TestDbSessionPropertyManagerMysql + extends TestDbSessionPropertyManager +{ + public TestDbSessionPropertyManagerMysql() throws Exception + { + super("mysql"); + } +} diff --git a/presto-delta/src/main/java/com/facebook/presto/delta/DeltaMetadata.java b/presto-delta/src/main/java/com/facebook/presto/delta/DeltaMetadata.java index d0b45969dd390..22e67e9d47f6e 100644 --- a/presto-delta/src/main/java/com/facebook/presto/delta/DeltaMetadata.java +++ b/presto-delta/src/main/java/com/facebook/presto/delta/DeltaMetadata.java @@ -231,7 +231,7 @@ public DeltaTableHandle getTableHandle(ConnectorSession session, SchemaTableName } @Override - public List getTableLayouts( + public ConnectorTableLayoutResult getTableLayoutForConstraint( ConnectorSession session, ConnectorTableHandle table, Constraint constraint, @@ -259,7 +259,7 @@ public List getTableLayouts( ImmutableList.of(), Optional.empty()); - return ImmutableList.of(new ConnectorTableLayoutResult(newLayout, unenforcedPredicate)); + return new ConnectorTableLayoutResult(newLayout, unenforcedPredicate); } @Override diff --git a/presto-docs/src/main/sphinx/admin/session-property-managers.rst b/presto-docs/src/main/sphinx/admin/session-property-managers.rst index 5ebdfb1a88bc3..4035e8dbb474f 100644 --- a/presto-docs/src/main/sphinx/admin/session-property-managers.rst +++ b/presto-docs/src/main/sphinx/admin/session-property-managers.rst @@ -5,10 +5,12 @@ Session Property Managers Administrators can add session properties to control the behavior for subsets of their workload. These properties are defaults and can be overridden by users (if authorized to do so). Session properties can be used to control resource usage, enable or disable features, and change query -characteristics. Session property managers are pluggable. +characteristics. Session property managers are pluggable. A session property manager can either +be database-based or file-based. For production environments, the database-based manager is +recommended as the properties can be updated without requiring a cluster restart. -Add an ``etc/session-property-config.properties`` file with the following contents to enable -the built-in manager that reads a JSON config file: +To enable a built-in manager that reads a JSON configuration file, add an +``etc/session-property-config.properties`` file with the following contents: .. code-block:: none @@ -24,6 +26,66 @@ by default. All matching rules contribute to constructing a list of session prop are applied in the order they are specified. Rules specified later in the file override values for properties that have been previously encountered. + +For the database-based built-in manager, add an +``etc/session-property-config.properties`` file with the following contents: + +.. code-block:: text + + session-property-config.configuration-manager=db + session-property-manager.db.url=jdbc:mysql://localhost:3306/session_properties?user=user&password=pass&createDatabaseIfNotExist=true + session-property-manager.db.refresh-period=50s + +Change the value of ``session-property-manager.db.url`` to the JDBC URL of a database. + +``session-property-manager.db.refresh-period`` should be set to how often Presto refreshes +to fetch the latest session properties from the database. + +This database consists of three tables: ``session_specs``, ``session_client_tags`` and ``session_property_values``. +Presto will create the database on startup if you set ``createDatabaseIfNotExist`` to ``true`` in your JDBC URL. +If the tables do not exist, Presto will create them on startup. + +.. code-block:: text + + mysql> DESCRIBE session_specs; + +-----------------------------+--------------+------+-----+---------+----------------+ + | Field | Type | Null | Key | Default | Extra | + +-----------------------------+--------------+------+-----+---------+----------------+ + | spec_id | bigint | NO | PRI | NULL | auto_increment | + | user_regex | varchar(512) | YES | | NULL | | + | source_regex | varchar(512) | YES | | NULL | | + | query_type | varchar(512) | YES | | NULL | | + | group_regex | varchar(512) | YES | | NULL | | + | client_info_regex | varchar(512) | YES | | NULL | | + | override_session_properties | tinyint(1) | YES | | NULL | | + | priority | int | NO | | NULL | | + +-----------------------------+--------------+------+-----+---------+----------------+ + 8 rows in set (0.016 sec) + +.. code-block:: text + + mysql> DESCRIBE session_client_tags; + +-------------+--------------+------+-----+---------+-------+ + | Field | Type | Null | Key | Default | Extra | + +-------------+--------------+------+-----+---------+-------+ + | tag_spec_id | bigint | NO | PRI | NULL | | + | client_tag | varchar(512) | NO | PRI | NULL | | + +-------------+--------------+------+-----+---------+-------+ + 2 rows in set (0.062 sec) + +.. code-block:: text + + mysql> DESCRIBE session_property_values; + +--------------------------+--------------+------+-----+---------+-------+ + | Field | Type | Null | Key | Default | Extra | + +--------------------------+--------------+------+-----+---------+-------+ + | property_spec_id | bigint | NO | PRI | NULL | | + | session_property_name | varchar(512) | NO | PRI | NULL | | + | session_property_value | varchar(512) | YES | | NULL | | + | session_property_catalog | varchar(512) | YES | | NULL | | + +--------------------------+--------------+------+-----+---------+-------+ + 3 rows in set (0.009 sec) + Match Rules ----------- @@ -52,9 +114,15 @@ Match Rules Note that once a session property has been overridden by ANY rule it remains overridden even if later higher precedence rules change the value, but don't specify override. -* ``sessionProperties``: map with string keys and values. Each entry is a system or catalog property name and +* ``sessionProperties``: map with string keys and values. Each entry is a system property name and corresponding value. Values must be specified as strings, no matter the actual data type. +* ``catalogSessionProperties``: map with string keys corresponding to the catalog name, and a map with string keys + and values as the value. Each entry is a catalog name and corresponding map of session property values. + +* For the database session property manager, catalog & system session properties are located in the same table. + ``session_property_catalog`` should be null for system session properties. + Example ------- @@ -71,6 +139,8 @@ Consider the following set of requirements: * All high memory ETL queries (tagged with 'high_mem_etl') are routed to subgroups under the ``global.pipeline`` group, and must be configured to enable :doc:`/admin/exchange-materialization`. +* All iceberg catalog queries should override the ``delete-as-join-rewrite-enabled`` property + These requirements can be expressed with the following rules: .. code-block:: json @@ -104,5 +174,12 @@ These requirements can be expressed with the following rules: "partitioning_provider_catalog": "hive", "hash_partition_count": 4096 } + }, + { + "catalogSessionProperties": { + "iceberg": { + "delete_as_join_rewrite_enabled": "true" + } + } } ] diff --git a/presto-docs/src/main/sphinx/conf.py b/presto-docs/src/main/sphinx/conf.py index 4f1e7342c5a16..ea71636098d47 100644 --- a/presto-docs/src/main/sphinx/conf.py +++ b/presto-docs/src/main/sphinx/conf.py @@ -135,6 +135,7 @@ def get_version(): 'features': [ 'toc.follow', 'toc.sticky', + 'content.code.copy', ], 'palette': [ { diff --git a/presto-docs/src/main/sphinx/connector/base-arrow-flight.rst b/presto-docs/src/main/sphinx/connector/base-arrow-flight.rst index e1676d3e05249..a68b4308a9d73 100644 --- a/presto-docs/src/main/sphinx/connector/base-arrow-flight.rst +++ b/presto-docs/src/main/sphinx/connector/base-arrow-flight.rst @@ -59,6 +59,36 @@ Property Name Description ``arrow-flight.server-ssl-enabled`` Port is ssl enabled ========================================== ============================================================== +Mutual TLS (mTLS) Support +------------------------- + + +To connect the Presto client to an Arrow Flight server with mutual TLS (mTLS) enabled, you must configure the client to present a valid certificate and key that the server can validate. This enhances security by ensuring both the client and server authenticate each other. + +To enable mTLS, the following properties must be configured: + +- ``arrow-flight.server-ssl-enabled=true``: Explicitly enables TLS for the connection. +- ``arrow-flight.server-ssl-certificate``: Path to the server's SSL certificate. +- ``arrow-flight.client-ssl-certificate``: Path to the client's SSL certificate. +- ``arrow-flight.client-ssl-key``: Path to the client's SSL private key. + +These properties must be used alongside the existing SSL configurations for the server, such as ``arrow-flight.server-ssl-certificate`` and ``arrow-flight.server-ssl-enabled=true``. Make sure the server is configured to trust the client certificates (typically via a shared CA). + +Below is an example code snippet to configure the Arrow Flight server with mTLS: + +.. code-block:: java + + File certChainFile = new File("src/test/resources/certs/server.crt"); + File privateKeyFile = new File("src/test/resources/certs/server.key"); + File caCertFile = new File("src/test/resources/certs/ca.crt"); + + server = FlightServer.builder(allocator, location, new TestingArrowProducer(allocator)) + .useTls(certChainFile, privateKeyFile) + .useMTlsClientVerification(caCertFile) + .build(); + + server.start(); + Querying Arrow-Flight --------------------- diff --git a/presto-docs/src/main/sphinx/connector/bigquery.rst b/presto-docs/src/main/sphinx/connector/bigquery.rst index d71c600b758c8..4e90e92555da0 100644 --- a/presto-docs/src/main/sphinx/connector/bigquery.rst +++ b/presto-docs/src/main/sphinx/connector/bigquery.rst @@ -137,6 +137,9 @@ Property Description ``bigquery.max-read-rows-retries`` The number of retries in case of retryable server issues ``3`` ``bigquery.credentials-key`` credentials key (base64 encoded) None. See `authentication <#authentication>`_ ``bigquery.credentials-file`` JSON credentials file path None. See `authentication <#authentication>`_ +``case-sensitive-name-matching`` Enable case sensitive identifier support for schema and table ``false`` + names for the connector. When disabled, names are matched + case-insensitively using lowercase normalization. ========================================= ============================================================== ============================================== Data Types diff --git a/presto-docs/src/main/sphinx/connector/deltalake.rst b/presto-docs/src/main/sphinx/connector/deltalake.rst index 96c7899e56972..3b44a60f1deee 100644 --- a/presto-docs/src/main/sphinx/connector/deltalake.rst +++ b/presto-docs/src/main/sphinx/connector/deltalake.rst @@ -133,3 +133,47 @@ in the table ``sales.apac.sales_data``. Above query drops the external table ``sales.apac.sales_data_new``. This only drops the metadata for the table. The referenced data directory is not deleted. + +Delta Lake to PrestoDB type mapping +----------------------------------- + +Map of Delta Lake types to the relevant PrestoDB types: + +.. list-table:: Delta Lake to PrestoDB type mapping + :widths: 50, 50 + :header-rows: 1 + + * - Delta Lake type + - PrestoDB type + * - ``BOOLEAN`` + - ``BOOLEAN`` + * - ``SMALLINT`` + - ``SMALLINT`` + * - ``TINYINT`` + - ``TINYINT`` + * - ``INT`` + - ``INTEGER`` + * - ``LONG`` + - ``BIGINT`` + * - ``FLOAT`` + - ``REAL`` + * - ``DOUBLE`` + - ``DOUBLE`` + * - ``DECIMAL`` + - ``DECIMAL`` + * - ``STRING`` + - ``VARCHAR`` + * - ``BINARY`` + - ``VARBINARY`` + * - ``DATE`` + - ``DATE`` + * - ``TIMESTAMP_NTZ`` + - ``TIMESTAMP`` + * - ``TIMESTAMP`` + - ``TIMESTAMP WITH TIME ZONE`` + * - ``ARRAY`` + - ``ARRAY`` + * - ``MAP`` + - ``MAP`` + * - ``STRUCT`` + - ``ROW`` diff --git a/presto-docs/src/main/sphinx/connector/hive.rst b/presto-docs/src/main/sphinx/connector/hive.rst index 7e427b9834b3b..cfdb309eeaf9f 100644 --- a/presto-docs/src/main/sphinx/connector/hive.rst +++ b/presto-docs/src/main/sphinx/connector/hive.rst @@ -224,7 +224,12 @@ Property Name Description CopyOnFirstWriteConfiguration, it can result in silent failures where critical configuration properties are not correctly propagated. + + ``hive.orc.use-column-names`` Enable accessing ORC columns by name in the ORC file ``false`` + metadata, instead of their ordinal position. Also toggleable + through the ``hive.orc_use_column_names`` session property. ======================================================== ============================================================ ============ + .. _constructor: https://github.com/apache/hadoop/blob/02a9190af5f8264e25966a80c8f9ea9bb6677899/hadoop-common-project/hadoop-common/src/main/java/org/apache/hadoop/conf/Configuration.java#L844-L875 Avro Configuration Properties diff --git a/presto-docs/src/main/sphinx/connector/iceberg.rst b/presto-docs/src/main/sphinx/connector/iceberg.rst index 2f8cec12f1c35..7b9adc3c0933a 100644 --- a/presto-docs/src/main/sphinx/connector/iceberg.rst +++ b/presto-docs/src/main/sphinx/connector/iceberg.rst @@ -93,6 +93,10 @@ Property Name Description ``iceberg.hive.table-refresh.backoff-scale-factor`` The multiple used to scale subsequent wait time between 4.0 retries. + +``iceberg.engine.hive.lock-enabled`` Whether to use locks to ensure atomicity of commits. true + This will turn off locks but is overridden at a table level + with the table configuration ``engine.hive.lock-enabled``. ======================================================== ============================================================= ============ Nessie catalog @@ -537,7 +541,8 @@ Property Name Description assign a split to. Splits which read data from the same file within the same chunk will hash to the same node. A smaller chunk size will result in a higher probability splits being distributed evenly across - the cluster, but reduce locality. + the cluster, but reduce locality. + See :ref:`develop/connectors:Node Selection Strategy`. ``iceberg.parquet_dereference_pushdown_enabled`` Overrides the behavior of the connector property Yes No ``iceberg.enable-parquet-dereference-pushdown`` in the current session. ===================================================== ======================================================================= =================== ============================================= diff --git a/presto-docs/src/main/sphinx/connector/mongodb.rst b/presto-docs/src/main/sphinx/connector/mongodb.rst index f40e7b0c50554..63a9bdedeb3af 100644 --- a/presto-docs/src/main/sphinx/connector/mongodb.rst +++ b/presto-docs/src/main/sphinx/connector/mongodb.rst @@ -51,6 +51,10 @@ Property Name Description ``mongodb.write-concern`` The write concern ``mongodb.required-replica-set`` The required replica set name ``mongodb.cursor-batch-size`` The number of elements to return in a batch +``case-sensitive-name-matching`` Enable case-sensitive identifier support for schema, + table, and column names for the connector. When disabled, + names are matched case-insensitively using lowercase + normalization. Default is ``false`` ===================================== ============================================================== ``mongodb.seeds`` diff --git a/presto-docs/src/main/sphinx/connector/postgresql.rst b/presto-docs/src/main/sphinx/connector/postgresql.rst index 256f7f56e72bb..065b54ec11baa 100644 --- a/presto-docs/src/main/sphinx/connector/postgresql.rst +++ b/presto-docs/src/main/sphinx/connector/postgresql.rst @@ -149,6 +149,7 @@ The connector maps PostgreSQL types to the corresponding PrestoDB types: - ``VARCHAR`` * - ``GEOGRAPHY`` - ``VARCHAR`` + No other types are supported. PrestoDB to PostgreSQL type mapping diff --git a/presto-docs/src/main/sphinx/connector/system.rst b/presto-docs/src/main/sphinx/connector/system.rst index 9f0b8038e4479..0e388c38627fc 100644 --- a/presto-docs/src/main/sphinx/connector/system.rst +++ b/presto-docs/src/main/sphinx/connector/system.rst @@ -36,7 +36,40 @@ System Connector Tables ``metadata.catalogs`` ^^^^^^^^^^^^^^^^^^^^^ -The catalogs table contains the list of available catalogs. +The catalogs table contains the list of available catalogs. The columns in ``metadata.catalogs`` are: + +======================================= ====================================================================== +Column Name Description +======================================= ====================================================================== +``catalog_name`` The value of this column is derived from the names of + catalog.properties files present under ``etc/catalog`` path under + presto installation directory. Everything except the suffix + ``.properties`` is treated as the catalog name. For example, if there + is a file named ``my_catalog.properties``, then ``my_catalog`` will be + listed as the value for this column. + +``connector_id`` The values in this column are a duplicate of the values in the + ``catalog_name`` column. + +``connector_name`` This column represents the actual name of the underlying connector + that a particular catalog is using. This column contains the value of + ``connector.name`` property from the catalog.properties file. +======================================= ====================================================================== + +Example: + +Suppose a user configures a single catalog by creating a file named ``my_catalog.properties`` with the +below contents:: + + connector.name=hive-hadoop2 + hive.metastore.uri=thrift://localhost:9083 + +``metadata.catalogs`` table will show below output:: + + presto> select * from system.metadata.catalogs; + catalog_name | connector_id | connector_name + --------------+--------------+---------------- + my_catalog | my_catalog | hive-hadoop2 ``metadata.schema_properties`` ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ diff --git a/presto-docs/src/main/sphinx/develop/connectors.rst b/presto-docs/src/main/sphinx/develop/connectors.rst index c77877890e3d3..24713c68cc176 100644 --- a/presto-docs/src/main/sphinx/develop/connectors.rst +++ b/presto-docs/src/main/sphinx/develop/connectors.rst @@ -8,7 +8,8 @@ you adapt your data source to the API expected by Presto, you can write queries against this data. ConnectorSplit ----------------- +-------------- + Instances of your connector splits. The ``getNodeSelectionStrategy`` method indicates the node affinity @@ -81,3 +82,14 @@ Given a split and a list of columns, the record set provider is responsible for delivering data to the Presto execution engine. It creates a ``RecordSet``, which in turn creates a ``RecordCursor`` that is used by Presto to read the column values for each row. + +Node Selection Strategy +----------------------- + +The node selection strategy is specified by a connector on each split. The possible values are: + +* HARD_AFFINITY - The Presto runtime must schedule this split on the nodes specified on ``ConnectorSplit#getPreferredNodes``. +* SOFT_AFFINITY - The Presto runtime should prefer ``ConnectorSplit#getPreferredNodes`` nodes, but doesn't have to. Use this value primarily for caching. +* NO_PREFERENCE - No preference. + +Use the ``node_selection_strategy`` session property in Hive and Iceberg to override this. \ No newline at end of file diff --git a/presto-docs/src/main/sphinx/functions/array.rst b/presto-docs/src/main/sphinx/functions/array.rst index 7cd3c60589dde..fde1e78f07581 100644 --- a/presto-docs/src/main/sphinx/functions/array.rst +++ b/presto-docs/src/main/sphinx/functions/array.rst @@ -198,6 +198,18 @@ Array Functions -1, IF(cardinality(x) = cardinality(y), 0, 1))); -- [[1, 2], [2, 3, 1], [4, 2, 1, 4]] +.. function:: array_sort(array(T), function(T,U)) -> array(T) + + Sorts and returns the ``array`` using a lambda function to extract sorting keys. The function is applied + to each element of the array to produce a key, and the array is sorted based on these keys in ascending order. + Null array elements and null keys are placed at the end. :: + + SELECT array_sort(ARRAY['pear', 'apple', 'banana', 'kiwi'], x -> length(x)); -- ['pear', 'kiwi', 'apple', 'banana'] + SELECT array_sort(ARRAY[5, 20, 3, 9, 100], x -> x); -- [3, 5, 9, 20, 100] + SELECT array_sort(ARRAY['apple', NULL, 'banana', NULL], x -> length(x)); -- ['apple', 'banana', NULL, NULL] + SELECT array_sort(ARRAY[CAST(0.0 AS DOUBLE), CAST('NaN' AS DOUBLE), CAST('Infinity' AS DOUBLE), CAST('-Infinity' AS DOUBLE)], x -> x); -- [-Infinity, 0.0, Infinity, NaN] + SELECT array_sort(ARRAY[ROW('a', 3), ROW('b', 1), ROW('c', 2)], x -> x[2]); -- [ROW('b', 1), ROW('c', 2), ROW('a', 3)] + .. function:: array_sort_desc(x) -> array Returns the ``array`` sorted in the descending order. Elements of the ``array`` must be orderable. @@ -207,6 +219,18 @@ Array Functions SELECT array_sort_desc(ARRAY [null, 100, null, 1, 10, 50]); -- [100, 50, 10, 1, null, null] SELECT array_sort_desc(ARRAY [ARRAY ["a", null], null, ARRAY ["a"]); -- [["a", null], ["a"], null] +.. function:: array_sort_desc(array(T), function(T,U)) -> array(T) + + Sorts and returns the ``array`` in descending order using a lambda function to extract sorting keys. + The function is applied to each element of the array to produce a key, and the array is sorted based + on these keys in descending order. Null array elements and null keys are placed at the end. :: + + SELECT array_sort_desc(ARRAY['pear', 'apple', 'banana', 'kiwi'], x -> length(x)); -- ['banana', 'apple', 'pear', 'kiwi'] + SELECT array_sort_desc(ARRAY[5, 20, 3, 9, 100], x -> x); -- [100, 20, 9, 5, 3] + SELECT array_sort_desc(ARRAY['apple', NULL, 'banana', NULL], x -> length(x)); -- ['banana', 'apple', NULL, NULL] + SELECT array_sort_desc(ARRAY[CAST(0.0 AS DOUBLE), CAST('NaN' AS DOUBLE), CAST('Infinity' AS DOUBLE), CAST('-Infinity' AS DOUBLE)], x -> x); -- [NaN, Infinity, 0.0, -Infinity] + SELECT array_sort_desc(ARRAY[ROW('a', 3), ROW('b', 1), ROW('c', 2)], x -> x[2]); -- [ROW('a', 3), ROW('c', 2), ROW('b', 1)] + .. function:: array_split_into_chunks(array(T), int) -> array(array(T)) Returns an ``array`` of arrays splitting the input ``array`` into chunks of given length. diff --git a/presto-docs/src/main/sphinx/installation/deploy-brew.rst b/presto-docs/src/main/sphinx/installation/deploy-brew.rst index 51cb3008da70e..bed202f4bd4d0 100644 --- a/presto-docs/src/main/sphinx/installation/deploy-brew.rst +++ b/presto-docs/src/main/sphinx/installation/deploy-brew.rst @@ -1,37 +1,31 @@ -===================================== -Deploy Presto on a Mac using Homebrew -===================================== +============================ +Deploy Presto using Homebrew +============================ -- If you are deploying Presto on an Intel Mac, see `Deploy Presto on an Intel Mac using Homebrew`_. +This guide explains how to install and get started with Presto on macOS, Linux or WSL2 using the Homebrew package manager. -- If you are deploying Presto on an Apple Silicon Mac that has an M1 or M2 chip, see `Deploy Presto on an Apple Silicon Mac using Homebrew`_. +Prerequisites +------------- -Deploy Presto on an Intel Mac using Homebrew --------------------------------------------- -*Note*: These steps were developed and tested on Mac OS X on Intel. These steps will not work with Apple Silicon (M1 or M2) chips. - -Following these steps, you will: - -- install the Presto service and CLI on an Intel Mac using `Homebrew `_ -- start and stop the Presto service -- start the Presto CLI +`Homebrew `_ installed. Install Presto -^^^^^^^^^^^^^^ +-------------- -Follow these steps to install Presto on an Intel Mac using `Homebrew `_. +Run the following command to install the latest version of Presto using the `Homebrew Formulae `_: -1. If you do not have brew installed, run the following command: +.. code-block:: none - ``/bin/bash -c "$(curl -fsSL https://raw.githubusercontent.com/Homebrew/install/master/install.sh)"`` + brew install prestodb -2. To install Presto, run the following command: +Homebrew installs packages in the ``Cellar`` directory, which can be found with this command: - ``brew install prestodb`` +.. code-block:: none - Presto is installed in the directory */usr/local/Cellar/prestodb/.* + brew --cellar -The following files are created in the *libexec/etc* directory in the Presto install directory: +The directory ``$(brew --cellar)/prestodb//libexec`` contains the Presto files used to run and configure the service. +For example, the ``etc`` directory within the Presto installation contains the following default configuration files: - node.properties - jvm.config @@ -39,35 +33,28 @@ The following files are created in the *libexec/etc* directory in the Presto ins - log.properties - catalog/jmx.properties -For example, the full path to the node.properties file is */usr/local/Cellar/prestodb//libexec/etc/node.properties*. - -The Presto CLI is installed in the *bin* directory of the Presto install directory: */usr/local/Cellar/prestodb//bin*. - -The executables are added to */usr/local/bin* path and should be available as part of $PATH. - Start and Stop Presto -^^^^^^^^^^^^^^^^^^^^^ +--------------------- -To start Presto, use the ``presto-server`` helper script. +Presto is installed with the ``presto-server`` helper script, which simplifies managing the cluster. +For example, run the following command to start the Presto service in the foreground: -To start the Presto service in the background, run the following command: - -``presto-server start`` +.. code-block:: none -To start the Presto service in the foreground, run the following command: + presto-server run -``presto-server run`` +To stop Presto from running in the foreground, press ``Ctrl + C`` until the terminal prompt appears, or close the terminal. -To stop the Presto service in the background, run the following command: +For more available commands and options, use help: -``presto-server stop`` +.. code-block:: none -To stop the Presto service in the foreground, close the terminal or select Ctrl + C until the terminal prompt is shown. + presto-server --help Open the Presto Console -^^^^^^^^^^^^^^^^^^^^^^^ +----------------------- -After starting Presto, you can access the web UI at the default port ``8080`` using the following link in a browser: +After starting the service, Presto Console can be accessible at the default port ``8080`` using the following link in a browser: .. code-block:: none @@ -79,117 +66,23 @@ After starting Presto, you can access the web UI at the default port ``8080`` us For more information about the Presto Console, see :doc:`/clients/presto-console`. Start the Presto CLI -^^^^^^^^^^^^^^^^^^^^ +-------------------- The Presto CLI is a terminal-based interactive shell for running queries, and is a `self-executing `_ JAR file that acts like a normal UNIX executable. -The Presto CLI is installed in the *bin* directory of the Presto install directory: */usr/local/Cellar/prestodb//bin*. - -To run the Presto CLI, run the following command: - -``presto`` - -The Presto CLI starts and displays the prompt ``presto>``. - -For more information, see :doc:`/clients/presto-cli`. - -Deploy Presto on an Apple Silicon Mac using Homebrew ----------------------------------------------------- -*Note*: These steps were developed and tested on Mac OS X on Apple Silicon. These steps will not work with Intel chips. - -Following these steps, you will: - -- install the Presto service and CLI on an Apple Silicon Mac using `Homebrew `_ -- start and stop the Presto service -- start the Presto CLI - -Install Presto -^^^^^^^^^^^^^^ - -Follow these steps to install Presto on an Apple Silicon Mac using `Homebrew `_. - -1. If you do not have brew installed, run the following command: - - ``arch -x86_64 /bin/bash -c "$(curl -fsSL https://raw.githubusercontent.com/Homebrew/install/master/install.sh)"`` - - This installs Homebrew into ``/usr/local/bin``. - - *Note*: The default installation of Homebrew on Apple Silicon is to ``/opt/homebrew``. - -2. To allow the shell to look for Homebrew in ``/usr/local/bin`` before it looks for Homebrew in ``/opt/homebrew``, run the following command: - - ``export PATH=/usr/local/bin:$PATH`` - -3. To install Presto, run the following command: - - ``arch -x86_64 brew install prestodb`` - - Presto is installed in the directory */usr/local/Cellar/prestodb/.* The executables ``presto`` - and ``presto-server`` are installed in ``/usr/local/bin/``. - -The following files are created in the *libexec/etc* directory in the Presto install directory: - -- node.properties -- jvm.config -- config.properties -- log.properties -- catalog/jmx.properties - -For example, the full path to the node.properties file is */usr/local/Cellar/prestodb//libexec/etc/node.properties*. - -The Presto CLI is installed in the *bin* directory of the Presto install directory: */usr/local/Cellar/prestodb//bin*. - -The executables are added to */usr/local/bin* path and should be available as part of $PATH. - -Start and Stop Presto -^^^^^^^^^^^^^^^^^^^^^ - -To start Presto, use the ``presto-server`` helper script. - -To start the Presto service in the background, run the following command: - -``arch -x86_64 presto-server start`` - -To start the Presto service in the foreground, run the following command: - -``arch -x86_64 presto-server run`` - -To stop the Presto service in the background, run the following command: - -``presto-server stop`` - -To stop the Presto service in the foreground, close the terminal or select Ctrl + C until the terminal prompt is shown. - -Open the Presto Console -^^^^^^^^^^^^^^^^^^^^^^^ - -After starting Presto, you can access the web UI at the default port ``8080`` using the following link in a browser: +The Presto CLI is installed in the directory ``$(brew --cellar)/prestodb//bin``. +To run the Presto CLI, use the following command: .. code-block:: none - http://localhost:8080 - -.. figure:: ../images/presto_console.png - :align: center - -For more information about the Presto Console, see :doc:`/clients/presto-console`. - -Start the Presto CLI -^^^^^^^^^^^^^^^^^^^^ - -The Presto CLI is a terminal-based interactive shell for running queries, and is a -`self-executing `_ -JAR file that acts like a normal UNIX executable. - -The Presto CLI is installed in the *bin* directory of the Presto install directory: */usr/local/Cellar/prestodb//bin*. -The executable ``presto`` is installed in ``/usr/local/bin/``. + presto -To run the Presto CLI, run the following command: +The Presto CLI starts and displays its prompt: -``presto`` +.. code-block:: none -The Presto CLI starts and displays the prompt ``presto>``. + presto> -For more information, see :doc:`/clients/presto-cli`. \ No newline at end of file +For more information, see :doc:`/clients/presto-cli`. diff --git a/presto-docs/src/main/sphinx/optimizer/history-based-optimization.rst b/presto-docs/src/main/sphinx/optimizer/history-based-optimization.rst index eeb5884e55be7..01e011889bbec 100644 --- a/presto-docs/src/main/sphinx/optimizer/history-based-optimization.rst +++ b/presto-docs/src/main/sphinx/optimizer/history-based-optimization.rst @@ -67,11 +67,11 @@ Session property Name Description ``restrict_history_based_optimization_to_complex_query`` Enable history based optimization only for complex queries, i.e. queries with join and aggregation. ``True`` ``history_input_table_statistics_matching_threshold`` Overrides the behavior of the configuration property ``hbo.history-matching-threshold`` ``hbo.history-matching-threshold`` in the current session. -``treat-low-confidence-zero-estimation-as-unknown`` Overrides the behavior of the configuration property +``treat_low_confidence_zero_estimation_unknown_enabled`` Overrides the behavior of the configuration property ``optimizer.treat-low-confidence-zero-estimation-as-unknown`` in the current session. ``optimizer.treat-low-confidence-zero-estimation-as-unknown`` -``confidence-based-broadcast`` Overrides the behavior of the configuration property +``confidence_based_broadcast_enabled`` Overrides the behavior of the configuration property ``optimizer.confidence-based-broadcast`` in the current session. ``optimizer.confidence-based-broadcast`` -``retry-query-with-history-based-optimization`` Overrides the behavior of the configuration property +``retry_query_with_history_based_optimization`` Overrides the behavior of the configuration property ``optimizer.retry-query-with-history-based-optimization`` in the current session. ``optimizer.retry-query-with-history-based-optimization`` =========================================================== ==================================================================================================== ============================================================== diff --git a/presto-docs/src/main/sphinx/plugin/native-sidecar-plugin.rst b/presto-docs/src/main/sphinx/plugin/native-sidecar-plugin.rst index e9cca5dc14438..760176afbe635 100644 --- a/presto-docs/src/main/sphinx/plugin/native-sidecar-plugin.rst +++ b/presto-docs/src/main/sphinx/plugin/native-sidecar-plugin.rst @@ -26,9 +26,12 @@ Property Name Description ``coordinator-sidecar-enabled`` Enables sidecar in the coordinator true ``native-execution-enabled`` Enables native execution true ``presto.default-namespace`` Sets the default function namespace `native.default` +``plugin.dir`` Specifies which directory under installation root `{root-directory}/native-plugins/` + to scan for plugins at startup. ============================================ ===================================================================== ============================== .. _sidecar-worker-properties: + Sidecar worker properties ^^^^^^^^^^^^^^^^^^^^^^^^^ Enable sidecar functionality with: @@ -68,7 +71,7 @@ Property Name Description ============================================ ===================================================================== ============================== Session properties ------------------ +------------------ These properties must be configured in ``etc/session-property-providers/native-worker.properties`` to use the session property provider of the ``NativeSidecarPlugin``. diff --git a/presto-docs/src/main/sphinx/presto_cpp/properties-session.rst b/presto-docs/src/main/sphinx/presto_cpp/properties-session.rst index 629040783be03..687dffbcd2f0a 100644 --- a/presto-docs/src/main/sphinx/presto_cpp/properties-session.rst +++ b/presto-docs/src/main/sphinx/presto_cpp/properties-session.rst @@ -97,6 +97,29 @@ If set to ``true``, disables the optimization in expression evaluation to delay This should only be used for debugging purposes. +``native_debug_memory_pool_name_regex`` +^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ + +* **Type:** ``varchar`` +* **Default value:** ``""`` + +Native Execution only. Regular expression pattern to match memory pool names for allocation callsite tracking. +Matched pools will also perform leak checks at destruction. Empty string disables tracking. + +This should only be used for debugging purposes. + +``native_debug_memory_pool_warn_threshold_bytes`` +^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ + +* **Type:** ``bigint`` +* **Default value:** ``0`` + +Native Execution only. Warning threshold for memory pool allocations. Logs callsites when exceeded. +Requires allocation tracking to be enabled with ``native_debug_memory_pool_name_regex``. +Accepts B/KB/MB/GB units. Set to 0B to disable. + +This should only be used for debugging purposes. + ``native_execution_type_rewrite_enabled`` ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ diff --git a/presto-docs/src/main/sphinx/presto_cpp/properties.rst b/presto-docs/src/main/sphinx/presto_cpp/properties.rst index cdb078a35a6da..2f1819029c21d 100644 --- a/presto-docs/src/main/sphinx/presto_cpp/properties.rst +++ b/presto-docs/src/main/sphinx/presto_cpp/properties.rst @@ -75,6 +75,7 @@ alphabetical order. - Custom type names are peeled in the coordinator. Only the actual base type is preserved. - ``CAST(col AS EnumType)`` is rewritten as ``CAST(col AS )``. - ``ENUM_KEY(EnumType)`` is rewritten as ``ELEMENT_AT(MAP(, VARCHAR))``. + This property can only be enabled with native execution. ``optimizer.optimize-hash-generation`` diff --git a/presto-docs/src/main/sphinx/release/release-0.294.rst b/presto-docs/src/main/sphinx/release/release-0.294.rst index 5fdc161bef1d0..73b7d871ff7de 100644 --- a/presto-docs/src/main/sphinx/release/release-0.294.rst +++ b/presto-docs/src/main/sphinx/release/release-0.294.rst @@ -12,6 +12,7 @@ Release 0.294 * Add mixed case support for schema and table names. `#24551 `_ * Add case-sensitive support for column names. It can be enabled for JDBC based connector by setting ``case-sensitive-name-matching=true`` at the catalog level. `#24983 `_ * Update ``presto-plan-checker-router-plugin router`` plugin to use ``EXPLAIN (TYPE VALIDATE)`` in place of ``EXPLAIN (TYPE DISTRIBUTED)``, enabling faster routing of queries to either native or Java clusters. `#25545 `_ +* From release 0.294, due to Maven Central publishing limitations, executable jar files including ``presto-cli``, ``presto-benchmark-driver``, and ``presto-test-server-launcher`` are no longer published in the Maven Central repository. These jars can now be found on the `Presto GitHub release page `_. **Details** =========== diff --git a/presto-docs/src/main/sphinx/security.rst b/presto-docs/src/main/sphinx/security.rst index f8c5432555e97..6d793c4c785e6 100644 --- a/presto-docs/src/main/sphinx/security.rst +++ b/presto-docs/src/main/sphinx/security.rst @@ -13,3 +13,4 @@ Security security/built-in-system-access-control security/internal-communication security/authorization + security/oauth2 diff --git a/presto-docs/src/main/sphinx/security/authorization.rst b/presto-docs/src/main/sphinx/security/authorization.rst index 6092a45ca6235..82649077ba3e3 100644 --- a/presto-docs/src/main/sphinx/security/authorization.rst +++ b/presto-docs/src/main/sphinx/security/authorization.rst @@ -46,7 +46,7 @@ so make sure you have authentication enabled. http-server.authentication.type=CERTIFICATE - It is also possible to specify other authentication types such as - ``KERBEROS``, ``PASSWORD`` and ``JWT``. Additional configuration may be + ``KERBEROS``, ``PASSWORD``, ``JWT``, and ``OAUTH2``. Additional configuration may be needed. .. code-block:: none diff --git a/presto-docs/src/main/sphinx/security/oauth2.rst b/presto-docs/src/main/sphinx/security/oauth2.rst new file mode 100644 index 0000000000000..65b676c42d394 --- /dev/null +++ b/presto-docs/src/main/sphinx/security/oauth2.rst @@ -0,0 +1,68 @@ +======================== +Oauth 2.0 Authentication +======================== + +Presto can be configured to enable frontend OAuth2 authentication over HTTPS for clients such as the CLI, JDBC, and ODBC drivers. OAuth2 provides a secure and flexible way to authenticate users by using an external identity provider (IdP), such as Okta, Auth0, Azure AD, or Google. + +OAuth2 authentication in Presto uses the Authorization Code Flow with PKCE and OpenID Connect (OIDC). The Presto coordinator initiates an OAuth2 challenge, and the client completes the flow by obtaining an access token from the identity provider. + +Presto Server Configuration +--------------------------- + +To enable OAuth2 authentication, configuration changes are made **only on the Presto coordinator**. No changes are required on the workers. + +Secure Communication +-------------------- + +Access to the Presto coordinator must be secured with HTTPS. You must configure a valid TLS certificate and keystore on the coordinator. See the `TLS setup guide `_ for details. + +OAuth2 Configuration +-------------------- + +Below are the key configuration properties for enabling OAuth2 authentication in ``config.properties``: + +.. code-block:: properties + + http-server.authentication.type=OAUTH2 + + http-server.authentication.oauth2.issuer=https://your-idp.com/oauth2/default + http-server.authentication.oauth2.client-id=your-client-id + http-server.authentication.oauth2.client-secret=your-client-secret + http-server.authentication.oauth2.scopes=openid,email,profile + http-server.authentication.oauth2.principal-field=sub + http-server.authentication.oauth2.groups-field=groups + http-server.authentication.oauth2.challenge-timeout=15m + http-server.authentication.oauth2.max-clock-skew=1m + http-server.authentication.oauth2.refresh-tokens=true + http-server.authentication.oauth2.oidc.discovery=true + http-server.authentication.oauth2.state-key=your-hmac-secret + http-server.authentication.oauth2.additional-audiences=your-client-id,another-audience + http-server.authentication.oauth2.user-mapping.pattern=(.*) + +It is worth noting that ``configuration-based-authorizer.role-regex-map.file-path`` must be configured if +authentication type is set to ``OAUTH2``. + +TLS Truststore for IdP +---------------------- + +If your IdP uses a custom or self-signed certificate, import it into the Java truststore on the Presto coordinator: + +.. code-block:: bash + + keytool -import \ + -keystore $JAVA_HOME/lib/security/cacerts \ + -trustcacerts \ + -alias idp_cert \ + -file idp_cert.crt + +Notes +----- + +- **Issuer**: The base URL of your IdP’s OIDC discovery endpoint. +- **Client ID/Secret**: Registered credentials for Presto in your IdP. +- **Scopes**: Must include ``openid``; others like ``email``, ``profile``, or ``groups`` are optional. +- **Principal Field**: The claim in the ID token used as the Presto username. +- **Groups Field**: Optional claim used for role-based access control. +- **State Key**: A secret used to sign the OAuth2 state parameter (HMAC). +- **Refresh Tokens**: Enable if your IdP supports issuing refresh tokens. +- **Callback**: When configuring your IdP the callback URI must be set to ``[presto]/oauth2/callback`` diff --git a/presto-druid/src/main/java/com/facebook/presto/druid/DruidMetadata.java b/presto-druid/src/main/java/com/facebook/presto/druid/DruidMetadata.java index bcc6df536e167..ce0d7f40e3ec7 100644 --- a/presto-druid/src/main/java/com/facebook/presto/druid/DruidMetadata.java +++ b/presto-druid/src/main/java/com/facebook/presto/druid/DruidMetadata.java @@ -85,11 +85,15 @@ public ConnectorTableHandle getTableHandle(ConnectorSession session, SchemaTable } @Override - public List getTableLayouts(ConnectorSession session, ConnectorTableHandle table, Constraint constraint, Optional> desiredColumns) + public ConnectorTableLayoutResult getTableLayoutForConstraint( + ConnectorSession session, + ConnectorTableHandle table, + Constraint constraint, + Optional> desiredColumns) { DruidTableHandle handle = (DruidTableHandle) table; ConnectorTableLayout layout = new ConnectorTableLayout(new DruidTableLayoutHandle(handle, constraint.getSummary())); - return ImmutableList.of(new ConnectorTableLayoutResult(layout, constraint.getSummary())); + return new ConnectorTableLayoutResult(layout, constraint.getSummary()); } @Override diff --git a/presto-elasticsearch/src/main/java/com/facebook/presto/elasticsearch/ElasticsearchMetadata.java b/presto-elasticsearch/src/main/java/com/facebook/presto/elasticsearch/ElasticsearchMetadata.java index aa8b62bbad357..e6b9b1c0e3ca7 100644 --- a/presto-elasticsearch/src/main/java/com/facebook/presto/elasticsearch/ElasticsearchMetadata.java +++ b/presto-elasticsearch/src/main/java/com/facebook/presto/elasticsearch/ElasticsearchMetadata.java @@ -158,11 +158,15 @@ public ElasticsearchTableHandle getTableHandle(ConnectorSession session, SchemaT } @Override - public List getTableLayouts(ConnectorSession session, ConnectorTableHandle table, Constraint constraint, Optional> desiredColumns) + public ConnectorTableLayoutResult getTableLayoutForConstraint( + ConnectorSession session, + ConnectorTableHandle table, + Constraint constraint, + Optional> desiredColumns) { ElasticsearchTableHandle handle = (ElasticsearchTableHandle) table; ConnectorTableLayout layout = new ConnectorTableLayout(new ElasticsearchTableLayoutHandle(handle, constraint.getSummary())); - return ImmutableList.of(new ConnectorTableLayoutResult(layout, constraint.getSummary())); + return new ConnectorTableLayoutResult(layout, constraint.getSummary()); } @Override diff --git a/presto-elasticsearch/src/test/java/com/facebook/presto/elasticsearch/ElasticsearchLoader.java b/presto-elasticsearch/src/test/java/com/facebook/presto/elasticsearch/ElasticsearchLoader.java index 5a0e736c70c44..3c7d62d7a7567 100644 --- a/presto-elasticsearch/src/test/java/com/facebook/presto/elasticsearch/ElasticsearchLoader.java +++ b/presto-elasticsearch/src/test/java/com/facebook/presto/elasticsearch/ElasticsearchLoader.java @@ -119,7 +119,7 @@ public void addResults(QueryStatusInfo statusInfo, QueryData data) } @Override - public Void build(Map setSessionProperties, Set resetSessionProperties) + public Void build(Map setSessionProperties, Set resetSessionProperties, String startTransactionId, boolean clearTransactionId) { return null; } diff --git a/presto-example-http/src/main/java/com/facebook/presto/example/ExampleMetadata.java b/presto-example-http/src/main/java/com/facebook/presto/example/ExampleMetadata.java index 9653a09277cac..9b6f4f2c504b9 100644 --- a/presto-example-http/src/main/java/com/facebook/presto/example/ExampleMetadata.java +++ b/presto-example-http/src/main/java/com/facebook/presto/example/ExampleMetadata.java @@ -81,11 +81,15 @@ public ExampleTableHandle getTableHandle(ConnectorSession session, SchemaTableNa } @Override - public List getTableLayouts(ConnectorSession session, ConnectorTableHandle table, Constraint constraint, Optional> desiredColumns) + public ConnectorTableLayoutResult getTableLayoutForConstraint( + ConnectorSession session, + ConnectorTableHandle table, + Constraint constraint, + Optional> desiredColumns) { ExampleTableHandle tableHandle = (ExampleTableHandle) table; ConnectorTableLayout layout = new ConnectorTableLayout(new ExampleTableLayoutHandle(tableHandle)); - return ImmutableList.of(new ConnectorTableLayoutResult(layout, constraint.getSummary())); + return new ConnectorTableLayoutResult(layout, constraint.getSummary()); } @Override diff --git a/presto-session-property-managers/pom.xml b/presto-file-session-property-manager/pom.xml similarity index 78% rename from presto-session-property-managers/pom.xml rename to presto-file-session-property-manager/pom.xml index 6386c244b92e5..5c3a4ceb2dddd 100644 --- a/presto-session-property-managers/pom.xml +++ b/presto-file-session-property-manager/pom.xml @@ -8,17 +8,21 @@ 0.295-SNAPSHOT - presto-session-property-managers - presto-session-property-managers - Presto - Session Property Managers + presto-file-session-property-manager + presto-file-session-property-manager + Presto - File Session Property Manager presto-plugin ${project.parent.basedir} - true + + com.facebook.presto + presto-session-property-managers-common + + com.facebook.airlift bootstrap @@ -88,12 +92,24 @@ provided + + io.airlift + slice + provided + + com.fasterxml.jackson.core jackson-annotations provided + + org.openjdk.jol + jol-core + provided + + com.facebook.presto @@ -112,5 +128,12 @@ testing test + + + com.facebook.presto + presto-session-property-managers-common + test-jar + test + diff --git a/presto-file-session-property-manager/src/main/java/com/facebook/presto/session/file/FileSessionPropertyConfigurationManagerPlugin.java b/presto-file-session-property-manager/src/main/java/com/facebook/presto/session/file/FileSessionPropertyConfigurationManagerPlugin.java new file mode 100644 index 0000000000000..573c62295572c --- /dev/null +++ b/presto-file-session-property-manager/src/main/java/com/facebook/presto/session/file/FileSessionPropertyConfigurationManagerPlugin.java @@ -0,0 +1,29 @@ +/* + * Licensed 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 com.facebook.presto.session.file; + +import com.facebook.presto.spi.Plugin; +import com.facebook.presto.spi.session.SessionPropertyConfigurationManagerFactory; +import com.google.common.collect.ImmutableList; + +public class FileSessionPropertyConfigurationManagerPlugin + implements Plugin +{ + @Override + public Iterable getSessionPropertyConfigurationManagerFactories() + { + return ImmutableList.of( + new FileSessionPropertyManagerFactory()); + } +} diff --git a/presto-session-property-managers/src/main/java/com/facebook/presto/session/FileSessionPropertyManager.java b/presto-file-session-property-manager/src/main/java/com/facebook/presto/session/file/FileSessionPropertyManager.java similarity index 60% rename from presto-session-property-managers/src/main/java/com/facebook/presto/session/FileSessionPropertyManager.java rename to presto-file-session-property-manager/src/main/java/com/facebook/presto/session/file/FileSessionPropertyManager.java index 9904baad66869..6f559a00f2b91 100644 --- a/presto-session-property-managers/src/main/java/com/facebook/presto/session/FileSessionPropertyManager.java +++ b/presto-file-session-property-manager/src/main/java/com/facebook/presto/session/file/FileSessionPropertyManager.java @@ -11,40 +11,36 @@ * See the License for the specific language governing permissions and * limitations under the License. */ -package com.facebook.presto.session; +package com.facebook.presto.session.file; import com.facebook.airlift.json.JsonCodec; import com.facebook.airlift.json.JsonCodecFactory; import com.facebook.airlift.json.JsonObjectMapperProvider; -import com.facebook.presto.spi.session.SessionConfigurationContext; -import com.facebook.presto.spi.session.SessionPropertyConfigurationManager; +import com.facebook.presto.session.AbstractSessionPropertyManager; +import com.facebook.presto.session.SessionMatchSpec; import com.fasterxml.jackson.databind.JsonMappingException; import com.fasterxml.jackson.databind.exc.UnrecognizedPropertyException; -import com.google.common.collect.ImmutableMap; +import com.google.common.collect.ImmutableList; import jakarta.inject.Inject; import java.io.IOException; import java.io.UncheckedIOException; import java.nio.file.Files; import java.nio.file.Path; -import java.util.HashMap; -import java.util.HashSet; import java.util.List; -import java.util.Map; -import java.util.Set; import static com.fasterxml.jackson.databind.DeserializationFeature.FAIL_ON_UNKNOWN_PROPERTIES; import static java.lang.String.format; import static java.util.Objects.requireNonNull; public class FileSessionPropertyManager - implements SessionPropertyConfigurationManager + extends AbstractSessionPropertyManager { public static final JsonCodec> CODEC = new JsonCodecFactory( () -> new JsonObjectMapperProvider().get().enable(FAIL_ON_UNKNOWN_PROPERTIES)) .listJsonCodec(SessionMatchSpec.class); - private final List sessionMatchSpecs; + private final ImmutableList sessionMatchSpecs; @Inject public FileSessionPropertyManager(FileSessionPropertyManagerConfig config) @@ -53,7 +49,7 @@ public FileSessionPropertyManager(FileSessionPropertyManagerConfig config) Path configurationFile = config.getConfigFile().toPath(); try { - sessionMatchSpecs = CODEC.fromJson(Files.readAllBytes(configurationFile)); + sessionMatchSpecs = ImmutableList.copyOf(CODEC.fromJson(Files.readAllBytes(configurationFile))); } catch (IOException e) { throw new UncheckedIOException(e); @@ -80,31 +76,8 @@ public FileSessionPropertyManager(FileSessionPropertyManagerConfig config) } @Override - public SystemSessionPropertyConfiguration getSystemSessionProperties(SessionConfigurationContext context) + protected List getSessionMatchSpecs() { - // later properties override earlier properties - Map defaultProperties = new HashMap<>(); - Set overridePropertyNames = new HashSet<>(); - for (SessionMatchSpec sessionMatchSpec : sessionMatchSpecs) { - Map newProperties = sessionMatchSpec.match(context); - defaultProperties.putAll(newProperties); - if (sessionMatchSpec.getOverrideSessionProperties().orElse(false)) { - overridePropertyNames.addAll(newProperties.keySet()); - } - } - - // Once a property has been overridden it stays that way and the value is updated by any rule - Map overrideProperties = new HashMap<>(); - for (String propertyName : overridePropertyNames) { - overrideProperties.put(propertyName, defaultProperties.get(propertyName)); - } - - return new SystemSessionPropertyConfiguration(ImmutableMap.copyOf(defaultProperties), ImmutableMap.copyOf(overrideProperties)); - } - - @Override - public Map> getCatalogSessionProperties(SessionConfigurationContext context) - { - return ImmutableMap.of(); + return sessionMatchSpecs; } } diff --git a/presto-session-property-managers/src/main/java/com/facebook/presto/session/FileSessionPropertyManagerConfig.java b/presto-file-session-property-manager/src/main/java/com/facebook/presto/session/file/FileSessionPropertyManagerConfig.java similarity index 96% rename from presto-session-property-managers/src/main/java/com/facebook/presto/session/FileSessionPropertyManagerConfig.java rename to presto-file-session-property-manager/src/main/java/com/facebook/presto/session/file/FileSessionPropertyManagerConfig.java index b1e72bed0aedd..54164b968b105 100644 --- a/presto-session-property-managers/src/main/java/com/facebook/presto/session/FileSessionPropertyManagerConfig.java +++ b/presto-file-session-property-manager/src/main/java/com/facebook/presto/session/file/FileSessionPropertyManagerConfig.java @@ -11,7 +11,7 @@ * See the License for the specific language governing permissions and * limitations under the License. */ -package com.facebook.presto.session; +package com.facebook.presto.session.file; import com.facebook.airlift.configuration.Config; import jakarta.validation.constraints.NotNull; diff --git a/presto-session-property-managers/src/main/java/com/facebook/presto/session/FileSessionPropertyManagerFactory.java b/presto-file-session-property-manager/src/main/java/com/facebook/presto/session/file/FileSessionPropertyManagerFactory.java similarity index 97% rename from presto-session-property-managers/src/main/java/com/facebook/presto/session/FileSessionPropertyManagerFactory.java rename to presto-file-session-property-manager/src/main/java/com/facebook/presto/session/file/FileSessionPropertyManagerFactory.java index aec39239cfbb5..08a49b74f54f5 100644 --- a/presto-session-property-managers/src/main/java/com/facebook/presto/session/FileSessionPropertyManagerFactory.java +++ b/presto-file-session-property-manager/src/main/java/com/facebook/presto/session/file/FileSessionPropertyManagerFactory.java @@ -11,7 +11,7 @@ * See the License for the specific language governing permissions and * limitations under the License. */ -package com.facebook.presto.session; +package com.facebook.presto.session.file; import com.facebook.airlift.bootstrap.Bootstrap; import com.facebook.airlift.json.JsonModule; diff --git a/presto-session-property-managers/src/main/java/com/facebook/presto/session/FileSessionPropertyManagerModule.java b/presto-file-session-property-manager/src/main/java/com/facebook/presto/session/file/FileSessionPropertyManagerModule.java similarity index 96% rename from presto-session-property-managers/src/main/java/com/facebook/presto/session/FileSessionPropertyManagerModule.java rename to presto-file-session-property-manager/src/main/java/com/facebook/presto/session/file/FileSessionPropertyManagerModule.java index 03fee8d5207dc..9d50e6144df4c 100644 --- a/presto-session-property-managers/src/main/java/com/facebook/presto/session/FileSessionPropertyManagerModule.java +++ b/presto-file-session-property-manager/src/main/java/com/facebook/presto/session/file/FileSessionPropertyManagerModule.java @@ -11,7 +11,7 @@ * See the License for the specific language governing permissions and * limitations under the License. */ -package com.facebook.presto.session; +package com.facebook.presto.session.file; import com.google.inject.Binder; import com.google.inject.Module; diff --git a/presto-file-session-property-manager/src/test/java/com/facebook/presto/session/file/TestFileSessionPropertyManager.java b/presto-file-session-property-manager/src/test/java/com/facebook/presto/session/file/TestFileSessionPropertyManager.java new file mode 100644 index 0000000000000..f071f2547b353 --- /dev/null +++ b/presto-file-session-property-manager/src/test/java/com/facebook/presto/session/file/TestFileSessionPropertyManager.java @@ -0,0 +1,120 @@ +/* + * Licensed 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 com.facebook.presto.session.file; + +import com.facebook.airlift.testing.TempFile; +import com.facebook.presto.session.AbstractTestSessionPropertyManager; +import com.facebook.presto.session.SessionMatchSpec; +import com.facebook.presto.spi.session.SessionPropertyConfigurationManager; +import com.google.common.collect.ImmutableMap; +import org.testng.annotations.Test; + +import java.io.IOException; +import java.nio.file.Files; +import java.nio.file.Path; +import java.util.Arrays; +import java.util.Map; +import java.util.Optional; + +import static com.facebook.presto.session.file.FileSessionPropertyManager.CODEC; +import static org.testng.Assert.assertEquals; + +public class TestFileSessionPropertyManager + extends AbstractTestSessionPropertyManager +{ + @Override + protected void assertProperties(Map defaultProperties, SessionMatchSpec... specs) + throws IOException + { + assertProperties(defaultProperties, ImmutableMap.of(), ImmutableMap.of(), specs); + } + + @Override + protected void assertProperties(Map defaultProperties, Map overrideProperties, SessionMatchSpec... specs) + throws IOException + { + assertProperties(defaultProperties, overrideProperties, ImmutableMap.of(), specs); + } + + protected void assertProperties(Map defaultProperties, Map overrideProperties, Map> catalogProperties, SessionMatchSpec... specs) + throws IOException + { + try (TempFile tempFile = new TempFile()) { + Path configurationFile = tempFile.path(); + Files.write(configurationFile, CODEC.toJsonBytes(Arrays.asList(specs))); + SessionPropertyConfigurationManager manager = new FileSessionPropertyManager(new FileSessionPropertyManagerConfig().setConfigFile(configurationFile.toFile())); + SessionPropertyConfigurationManager.SystemSessionPropertyConfiguration propertyConfiguration = manager.getSystemSessionProperties(CONTEXT); + assertEquals(propertyConfiguration.systemPropertyDefaults, defaultProperties); + assertEquals(propertyConfiguration.systemPropertyOverrides, overrideProperties); + assertEquals(manager.getCatalogSessionProperties(CONTEXT), catalogProperties); + } + } + + @Test + public void testNullSessionProperties() + throws IOException + { + ImmutableMap> catalogProperties = ImmutableMap.of("CATALOG", ImmutableMap.of("PROPERTY", "VALUE")); + SessionMatchSpec spec = new SessionMatchSpec( + Optional.empty(), + Optional.empty(), + Optional.empty(), + Optional.empty(), + Optional.empty(), + Optional.empty(), + Optional.empty(), + null, + catalogProperties); + + assertProperties(ImmutableMap.of(), ImmutableMap.of(), catalogProperties, spec); + } + + @Test + public void testNullCatalogSessionProperties() + throws IOException + { + Map properties = ImmutableMap.of("PROPERTY1", "VALUE1", "PROPERTY2", "VALUE2"); + SessionMatchSpec spec = new SessionMatchSpec( + Optional.empty(), + Optional.empty(), + Optional.empty(), + Optional.empty(), + Optional.empty(), + Optional.empty(), + Optional.empty(), + properties, + null); + + assertProperties(properties, spec); + } + + @Test(expectedExceptions = IllegalArgumentException.class, expectedExceptionsMessageRegExp = "Either sessionProperties or catalogSessionProperties must be provided") + public void testNullBothSessionProperties() + throws IOException + { + SessionMatchSpec spec = new SessionMatchSpec( + Optional.empty(), + Optional.empty(), + Optional.empty(), + Optional.empty(), + Optional.empty(), + Optional.empty(), + Optional.empty(), + null, + null); + + assertProperties(ImmutableMap.of(), spec); + } +} diff --git a/presto-session-property-managers/src/test/java/com/facebook/presto/session/TestFileSessionPropertyManagerConfig.java b/presto-file-session-property-manager/src/test/java/com/facebook/presto/session/file/TestFileSessionPropertyManagerConfig.java similarity index 92% rename from presto-session-property-managers/src/test/java/com/facebook/presto/session/TestFileSessionPropertyManagerConfig.java rename to presto-file-session-property-manager/src/test/java/com/facebook/presto/session/file/TestFileSessionPropertyManagerConfig.java index b555593fa09a0..43403ef267bce 100644 --- a/presto-session-property-managers/src/test/java/com/facebook/presto/session/TestFileSessionPropertyManagerConfig.java +++ b/presto-file-session-property-manager/src/test/java/com/facebook/presto/session/file/TestFileSessionPropertyManagerConfig.java @@ -11,12 +11,12 @@ * See the License for the specific language governing permissions and * limitations under the License. */ -package com.facebook.presto.session; +package com.facebook.presto.session.file; import com.google.common.collect.ImmutableMap; import org.testng.annotations.Test; -import java.io.File; +import java.nio.file.Paths; import java.util.Map; import static com.facebook.airlift.configuration.testing.ConfigAssertions.assertFullMapping; @@ -40,7 +40,7 @@ public void testExplicitPropertyMappings() .build(); FileSessionPropertyManagerConfig expected = new FileSessionPropertyManagerConfig() - .setConfigFile(new File("/test.json")); + .setConfigFile(Paths.get("/test.json").toFile()); assertFullMapping(properties, expected); } diff --git a/presto-function-namespace-managers/pom.xml b/presto-function-namespace-managers/pom.xml index 7e3a3e95425df..a882b50e35df6 100644 --- a/presto-function-namespace-managers/pom.xml +++ b/presto-function-namespace-managers/pom.xml @@ -148,7 +148,6 @@ org.mariadb.jdbc mariadb-java-client - 3.5.4 runtime diff --git a/presto-google-sheets/src/main/java/com/facebook/presto/google/sheets/SheetsMetadata.java b/presto-google-sheets/src/main/java/com/facebook/presto/google/sheets/SheetsMetadata.java index c833915fb6eda..c077c77c0e323 100644 --- a/presto-google-sheets/src/main/java/com/facebook/presto/google/sheets/SheetsMetadata.java +++ b/presto-google-sheets/src/main/java/com/facebook/presto/google/sheets/SheetsMetadata.java @@ -81,13 +81,15 @@ public SheetsTableHandle getTableHandle(ConnectorSession session, SchemaTableNam } @Override - public List getTableLayouts( - ConnectorSession session, ConnectorTableHandle table, - Constraint constraint, Optional> desiredColumns) + public ConnectorTableLayoutResult getTableLayoutForConstraint( + ConnectorSession session, + ConnectorTableHandle table, + Constraint constraint, + Optional> desiredColumns) { SheetsTableHandle tableHandle = (SheetsTableHandle) table; ConnectorTableLayout layout = new ConnectorTableLayout(new SheetsTableLayoutHandle(tableHandle)); - return ImmutableList.of(new ConnectorTableLayoutResult(layout, constraint.getSummary())); + return new ConnectorTableLayoutResult(layout, constraint.getSummary()); } @Override diff --git a/presto-hive-common/src/main/java/com/facebook/presto/hive/HiveCommonClientConfig.java b/presto-hive-common/src/main/java/com/facebook/presto/hive/HiveCommonClientConfig.java index 0b5a34b2c8ed3..1e54401247e2f 100644 --- a/presto-hive-common/src/main/java/com/facebook/presto/hive/HiveCommonClientConfig.java +++ b/presto-hive-common/src/main/java/com/facebook/presto/hive/HiveCommonClientConfig.java @@ -37,6 +37,7 @@ public class HiveCommonClientConfig private DataSize orcStreamBufferSize = new DataSize(8, MEGABYTE); private OrcWriteValidationMode orcWriterValidationMode = OrcWriteValidationMode.BOTH; private double orcWriterValidationPercentage; + private boolean useOrcColumnNames; private DataSize orcTinyStripeThreshold = new DataSize(8, MEGABYTE); private boolean parquetBatchReadOptimizationEnabled; private boolean parquetEnableBatchReaderVerification; @@ -183,6 +184,19 @@ public HiveCommonClientConfig setOrcWriterValidationPercentage(double orcWriterV return this; } + public boolean isUseOrcColumnNames() + { + return useOrcColumnNames; + } + + @Config("hive.orc.use-column-names") + @ConfigDescription("Access ORC columns using names from the file first, and fallback to Hive schema column names if not found to ensure backward compatibility with old data") + public HiveCommonClientConfig setUseOrcColumnNames(boolean useOrcColumnNames) + { + this.useOrcColumnNames = useOrcColumnNames; + return this; + } + @NotNull public DataSize getOrcTinyStripeThreshold() { diff --git a/presto-hive-common/src/main/java/com/facebook/presto/hive/HiveCommonSessionProperties.java b/presto-hive-common/src/main/java/com/facebook/presto/hive/HiveCommonSessionProperties.java index de4bce0e2b661..b505cd37d8188 100644 --- a/presto-hive-common/src/main/java/com/facebook/presto/hive/HiveCommonSessionProperties.java +++ b/presto-hive-common/src/main/java/com/facebook/presto/hive/HiveCommonSessionProperties.java @@ -42,6 +42,8 @@ public class HiveCommonSessionProperties public static final String RANGE_FILTERS_ON_SUBSCRIPTS_ENABLED = "range_filters_on_subscripts_enabled"; @VisibleForTesting public static final String PARQUET_BATCH_READ_OPTIMIZATION_ENABLED = "parquet_batch_read_optimization_enabled"; + @VisibleForTesting + public static final String ORC_USE_COLUMN_NAMES = "orc_use_column_names"; public static final String NODE_SELECTION_STRATEGY = "node_selection_strategy"; private static final String ORC_BLOOM_FILTERS_ENABLED = "orc_bloom_filters_enabled"; @@ -153,6 +155,11 @@ public HiveCommonSessionProperties(HiveCommonClientConfig hiveCommonClientConfig "use JNI based zstd decompression for reading ORC files", hiveCommonClientConfig.isZstdJniDecompressionEnabled(), true), + booleanProperty( + ORC_USE_COLUMN_NAMES, + "Access ORC columns using names from the file first, and fallback to Hive schema column names if not found to ensure backward compatibility with old data", + hiveCommonClientConfig.isUseOrcColumnNames(), + false), booleanProperty( PARQUET_BATCH_READ_OPTIMIZATION_ENABLED, "Is Parquet batch read optimization enabled", @@ -262,6 +269,11 @@ public static boolean isOrcZstdJniDecompressionEnabled(ConnectorSession session) return session.getProperty(ORC_ZSTD_JNI_DECOMPRESSION_ENABLED, Boolean.class); } + public static boolean isUseOrcColumnNames(ConnectorSession session) + { + return session.getProperty(ORC_USE_COLUMN_NAMES, Boolean.class); + } + public static boolean isParquetBatchReadsEnabled(ConnectorSession session) { return session.getProperty(PARQUET_BATCH_READ_OPTIMIZATION_ENABLED, Boolean.class); diff --git a/presto-hive-common/src/test/java/com/facebook/presto/hive/TestHiveCommonClientConfig.java b/presto-hive-common/src/test/java/com/facebook/presto/hive/TestHiveCommonClientConfig.java index da24f4e64bc32..42579c3048b92 100644 --- a/presto-hive-common/src/test/java/com/facebook/presto/hive/TestHiveCommonClientConfig.java +++ b/presto-hive-common/src/test/java/com/facebook/presto/hive/TestHiveCommonClientConfig.java @@ -45,6 +45,7 @@ public void testDefaults() .setOrcOptimizedWriterEnabled(true) .setOrcWriterValidationPercentage(0.0) .setOrcWriterValidationMode(OrcWriteValidation.OrcWriteValidationMode.BOTH) + .setUseOrcColumnNames(false) .setZstdJniDecompressionEnabled(false) .setParquetBatchReaderVerificationEnabled(false) .setParquetBatchReadOptimizationEnabled(false) @@ -71,6 +72,7 @@ public void testExplicitPropertyMappings() .put("hive.orc.optimized-writer.enabled", "false") .put("hive.orc.writer.validation-percentage", "0.16") .put("hive.orc.writer.validation-mode", "DETAILED") + .put("hive.orc.use-column-names", "true") .put("hive.zstd-jni-decompression-enabled", "true") .put("hive.enable-parquet-batch-reader-verification", "true") .put("hive.parquet-batch-read-optimization-enabled", "true") @@ -94,6 +96,7 @@ public void testExplicitPropertyMappings() .setOrcOptimizedWriterEnabled(false) .setOrcWriterValidationPercentage(0.16) .setOrcWriterValidationMode(OrcWriteValidation.OrcWriteValidationMode.DETAILED) + .setUseOrcColumnNames(true) .setZstdJniDecompressionEnabled(true) .setParquetBatchReaderVerificationEnabled(true) .setParquetBatchReadOptimizationEnabled(true) diff --git a/presto-hive-metastore/src/main/java/com/facebook/presto/hive/metastore/ExtendedHiveMetastore.java b/presto-hive-metastore/src/main/java/com/facebook/presto/hive/metastore/ExtendedHiveMetastore.java index dc972904a0ac4..4cfdc0d799a86 100644 --- a/presto-hive-metastore/src/main/java/com/facebook/presto/hive/metastore/ExtendedHiveMetastore.java +++ b/presto-hive-metastore/src/main/java/com/facebook/presto/hive/metastore/ExtendedHiveMetastore.java @@ -31,6 +31,7 @@ import java.util.Optional; import java.util.Set; import java.util.function.Function; +import java.util.function.Supplier; public interface ExtendedHiveMetastore { @@ -89,6 +90,8 @@ default void dropTableFromMetastore(MetastoreContext metastoreContext, String da */ MetastoreOperationResult replaceTable(MetastoreContext metastoreContext, String databaseName, String tableName, Table newTable, PrincipalPrivileges principalPrivileges); + MetastoreOperationResult persistTable(MetastoreContext metastoreContext, String databaseName, String tableName, Table newTable, PrincipalPrivileges principalPrivileges, Supplier update, Map additionalParameters); + MetastoreOperationResult renameTable(MetastoreContext metastoreContext, String databaseName, String tableName, String newDatabaseName, String newTableName); MetastoreOperationResult addColumn(MetastoreContext metastoreContext, String databaseName, String tableName, String columnName, HiveType columnType, String columnComment); diff --git a/presto-hive-metastore/src/main/java/com/facebook/presto/hive/metastore/InMemoryCachingHiveMetastore.java b/presto-hive-metastore/src/main/java/com/facebook/presto/hive/metastore/InMemoryCachingHiveMetastore.java index f281c0f600c3b..fdf142b4a25ef 100644 --- a/presto-hive-metastore/src/main/java/com/facebook/presto/hive/metastore/InMemoryCachingHiveMetastore.java +++ b/presto-hive-metastore/src/main/java/com/facebook/presto/hive/metastore/InMemoryCachingHiveMetastore.java @@ -49,6 +49,7 @@ import java.util.concurrent.ThreadLocalRandom; import java.util.function.Function; import java.util.function.Predicate; +import java.util.function.Supplier; import static com.facebook.presto.hive.HiveErrorCode.HIVE_CORRUPTED_PARTITION_CACHE; import static com.facebook.presto.hive.HiveErrorCode.HIVE_PARTITION_DROPPED_DURING_QUERY; @@ -450,6 +451,18 @@ private Map, PartitionStatistics> loadPartition return result.build(); } + @Override + public MetastoreOperationResult persistTable(MetastoreContext metastoreContext, String databaseName, String tableName, Table newTable, PrincipalPrivileges principalPrivileges, Supplier update, Map additionalParameters) + { + try { + return getDelegate().persistTable(metastoreContext, databaseName, tableName, newTable, principalPrivileges, update, additionalParameters); + } + finally { + invalidateTableCache(databaseName, tableName); + invalidateTableCache(newTable.getDatabaseName(), newTable.getTableName()); + } + } + @Override public void updateTableStatistics(MetastoreContext metastoreContext, String databaseName, String tableName, Function update) { diff --git a/presto-hive-metastore/src/main/java/com/facebook/presto/hive/metastore/RecordingHiveMetastore.java b/presto-hive-metastore/src/main/java/com/facebook/presto/hive/metastore/RecordingHiveMetastore.java index 6163ae58b8d55..4c70b9005690c 100644 --- a/presto-hive-metastore/src/main/java/com/facebook/presto/hive/metastore/RecordingHiveMetastore.java +++ b/presto-hive-metastore/src/main/java/com/facebook/presto/hive/metastore/RecordingHiveMetastore.java @@ -315,6 +315,13 @@ public MetastoreOperationResult replaceTable(MetastoreContext metastoreContext, return delegate.replaceTable(metastoreContext, databaseName, tableName, newTable, principalPrivileges); } + @Override + public MetastoreOperationResult persistTable(MetastoreContext metastoreContext, String databaseName, String tableName, Table newTable, PrincipalPrivileges principalPrivileges, Supplier update, Map additionalParameters) + { + verifyRecordingMode(); + return delegate.persistTable(metastoreContext, databaseName, tableName, newTable, principalPrivileges, update, additionalParameters); + } + @Override public MetastoreOperationResult renameTable(MetastoreContext metastoreContext, String databaseName, String tableName, String newDatabaseName, String newTableName) { diff --git a/presto-hive-metastore/src/main/java/com/facebook/presto/hive/metastore/file/FileHiveMetastore.java b/presto-hive-metastore/src/main/java/com/facebook/presto/hive/metastore/file/FileHiveMetastore.java index 0d2953f65046d..ee1879b73e698 100644 --- a/presto-hive-metastore/src/main/java/com/facebook/presto/hive/metastore/file/FileHiveMetastore.java +++ b/presto-hive-metastore/src/main/java/com/facebook/presto/hive/metastore/file/FileHiveMetastore.java @@ -86,6 +86,7 @@ import java.util.Optional; import java.util.Set; import java.util.function.Function; +import java.util.function.Supplier; import static com.facebook.presto.hive.HiveErrorCode.HIVE_METASTORE_ERROR; import static com.facebook.presto.hive.HiveErrorCode.HIVE_PARTITION_DROPPED_DURING_QUERY; @@ -134,6 +135,7 @@ public class FileHiveMetastore protected final HdfsEnvironment hdfsEnvironment; protected final HdfsContext hdfsContext; + protected final FileSystem metadataFileSystem; private final Path catalogDirectory; @@ -477,6 +479,44 @@ public synchronized MetastoreOperationResult replaceTable(MetastoreContext metas return EMPTY_RESULT; } + @Override + public synchronized MetastoreOperationResult persistTable( + MetastoreContext metastoreContext, + String databaseName, + String tableName, + Table newTable, + PrincipalPrivileges principalPrivileges, + Supplier update, + Map additionalParameters) + { + checkArgument(!newTable.getTableType().equals(TEMPORARY_TABLE), "temporary tables must never be stored in the metastore"); + + Table oldTable = getRequiredTable(metastoreContext, databaseName, tableName); + validateReplaceTableType(oldTable, newTable); + if (!oldTable.getDatabaseName().equals(databaseName) || !oldTable.getTableName().equals(tableName)) { + throw new PrestoException(HIVE_METASTORE_ERROR, "Replacement table must have same name"); + } + + Path tableMetadataDirectory = getTableMetadataDirectory(oldTable); + + deleteTablePrivileges(oldTable); + for (Entry> entry : principalPrivileges.getUserPrivileges().asMap().entrySet()) { + setTablePrivileges(metastoreContext, new PrestoPrincipal(USER, entry.getKey()), databaseName, tableName, entry.getValue()); + } + for (Entry> entry : principalPrivileges.getRolePrivileges().asMap().entrySet()) { + setTablePrivileges(metastoreContext, new PrestoPrincipal(ROLE, entry.getKey()), databaseName, tableName, entry.getValue()); + } + PartitionStatistics updatedStatistics = update.get(); + + TableMetadata updatedMetadata = new TableMetadata(newTable) + .withParameters(updateStatisticsParameters(newTable.getParameters(), updatedStatistics.getBasicStatistics())) + .withColumnStatistics(updatedStatistics.getColumnStatistics()); + + writeSchemaFile("table", tableMetadataDirectory, tableCodec, updatedMetadata, true); + + return EMPTY_RESULT; + } + @Override public synchronized MetastoreOperationResult renameTable(MetastoreContext metastoreContext, String databaseName, String tableName, String newDatabaseName, String newTableName) { @@ -1292,7 +1332,6 @@ private List getChildSchemaDirectories(Path metadataDirectory) if (!metadataFileSystem.isDirectory(metadataDirectory)) { return ImmutableList.of(); } - ImmutableList.Builder childSchemaDirectories = ImmutableList.builder(); for (FileStatus child : metadataFileSystem.listStatus(metadataDirectory)) { if (!child.isDirectory()) { diff --git a/presto-hive-metastore/src/main/java/com/facebook/presto/hive/metastore/glue/GlueHiveMetastore.java b/presto-hive-metastore/src/main/java/com/facebook/presto/hive/metastore/glue/GlueHiveMetastore.java index 20876e451db2a..ef782cd889fac 100644 --- a/presto-hive-metastore/src/main/java/com/facebook/presto/hive/metastore/glue/GlueHiveMetastore.java +++ b/presto-hive-metastore/src/main/java/com/facebook/presto/hive/metastore/glue/GlueHiveMetastore.java @@ -116,6 +116,7 @@ import java.util.concurrent.ExecutorCompletionService; import java.util.concurrent.Future; import java.util.function.Function; +import java.util.function.Supplier; import static com.facebook.presto.hive.HiveErrorCode.HIVE_METASTORE_ERROR; import static com.facebook.presto.hive.HiveErrorCode.HIVE_PARTITION_DROPPED_DURING_QUERY; @@ -589,6 +590,30 @@ public MetastoreOperationResult replaceTable(MetastoreContext metastoreContext, } } + public MetastoreOperationResult persistTable(MetastoreContext metastoreContext, String databaseName, String tableName, Table newTable, PrincipalPrivileges principalPrivileges, Supplier update, Map additionalParameters) + { + PartitionStatistics updatedStatistics = update.get(); + if (!updatedStatistics.getColumnStatistics().isEmpty()) { + throw new PrestoException(NOT_SUPPORTED, "Glue metastore does not support column level statistics"); + } + try { + TableInput newTableInput = GlueInputConverter.convertTable(newTable); + newTableInput.setParameters(updateStatisticsParameters(newTableInput.getParameters(), updatedStatistics.getBasicStatistics())); + stats.getUpdateTable().record(() -> glueClient.updateTable(new UpdateTableRequest() + .withCatalogId(catalogId) + .withDatabaseName(databaseName) + .withTableInput(newTableInput))); + + return EMPTY_RESULT; + } + catch (EntityNotFoundException e) { + throw new TableNotFoundException(new SchemaTableName(databaseName, tableName)); + } + catch (AmazonServiceException e) { + throw new PrestoException(HIVE_METASTORE_ERROR, e); + } + } + @Override public MetastoreOperationResult renameTable(MetastoreContext metastoreContext, String databaseName, String tableName, String newDatabaseName, String newTableName) { diff --git a/presto-hive-metastore/src/main/java/com/facebook/presto/hive/metastore/thrift/BridgingHiveMetastore.java b/presto-hive-metastore/src/main/java/com/facebook/presto/hive/metastore/thrift/BridgingHiveMetastore.java index d80d5f16570ef..e3787b7ebe91a 100644 --- a/presto-hive-metastore/src/main/java/com/facebook/presto/hive/metastore/thrift/BridgingHiveMetastore.java +++ b/presto-hive-metastore/src/main/java/com/facebook/presto/hive/metastore/thrift/BridgingHiveMetastore.java @@ -45,7 +45,10 @@ import com.facebook.presto.spi.statistics.ColumnStatisticType; import com.google.common.collect.ImmutableList; import com.google.common.collect.ImmutableMap; +import com.google.common.collect.Maps; import jakarta.inject.Inject; +import org.apache.hadoop.hive.common.StatsSetupConst; +import org.apache.hadoop.hive.metastore.api.EnvironmentContext; import org.apache.hadoop.hive.metastore.api.FieldSchema; import java.util.List; @@ -53,6 +56,7 @@ import java.util.Optional; import java.util.Set; import java.util.function.Function; +import java.util.function.Supplier; import java.util.stream.Collectors; import static com.facebook.presto.hive.metastore.MetastoreUtil.getPartitionNamesWithEmptyVersion; @@ -276,6 +280,21 @@ private MetastoreOperationResult alterTable(MetastoreContext metastoreContext, S return delegate.alterTable(metastoreContext, databaseName, tableName, table); } + private MetastoreOperationResult alterTableWithEnvironmentContext(MetastoreContext metastoreContext, String databaseName, String tableName, org.apache.hadoop.hive.metastore.api.Table table, EnvironmentContext environmentContext) + { + return delegate.alterTableWithEnvironmentContext(metastoreContext, databaseName, tableName, table, environmentContext); + } + + @Override + public MetastoreOperationResult persistTable(MetastoreContext metastoreContext, String databaseName, String tableName, Table newTable, PrincipalPrivileges principalPrivileges, Supplier update, Map additionalParameters) + { + checkArgument(!newTable.getTableType().equals(TEMPORARY_TABLE), "temporary tables must never be stored in the metastore"); + Map env = Maps.newHashMapWithExpectedSize(additionalParameters.size() + 1); + env.putAll(additionalParameters); + env.put(StatsSetupConst.DO_NOT_UPDATE_STATS, StatsSetupConst.TRUE); + return alterTableWithEnvironmentContext(metastoreContext, databaseName, tableName, toMetastoreApiTable(newTable, principalPrivileges, metastoreContext.getColumnConverter()), new EnvironmentContext(env)); + } + @Override public Optional getPartition(MetastoreContext metastoreContext, String databaseName, String tableName, List partitionValues) { diff --git a/presto-hive-metastore/src/main/java/com/facebook/presto/hive/metastore/thrift/HiveMetastore.java b/presto-hive-metastore/src/main/java/com/facebook/presto/hive/metastore/thrift/HiveMetastore.java index 0d499b18eca8d..f66484a7bdcba 100644 --- a/presto-hive-metastore/src/main/java/com/facebook/presto/hive/metastore/thrift/HiveMetastore.java +++ b/presto-hive-metastore/src/main/java/com/facebook/presto/hive/metastore/thrift/HiveMetastore.java @@ -36,6 +36,7 @@ import com.facebook.presto.spi.statistics.ColumnStatisticType; import com.google.common.collect.ImmutableList; import org.apache.hadoop.hive.metastore.api.Database; +import org.apache.hadoop.hive.metastore.api.EnvironmentContext; import org.apache.hadoop.hive.metastore.api.FieldSchema; import org.apache.hadoop.hive.metastore.api.Partition; import org.apache.hadoop.hive.metastore.api.Table; @@ -62,6 +63,8 @@ public interface HiveMetastore MetastoreOperationResult alterTable(MetastoreContext metastoreContext, String databaseName, String tableName, Table table); + MetastoreOperationResult alterTableWithEnvironmentContext(MetastoreContext metastoreContext, String databaseName, String tableName, Table table, EnvironmentContext environmentContext); + default List getDatabases(MetastoreContext metastoreContext, String pattern) { return getAllDatabases(metastoreContext); diff --git a/presto-hive-metastore/src/main/java/com/facebook/presto/hive/metastore/thrift/HiveMetastoreClient.java b/presto-hive-metastore/src/main/java/com/facebook/presto/hive/metastore/thrift/HiveMetastoreClient.java index 4880eedcdd6ee..caf27aeed8619 100644 --- a/presto-hive-metastore/src/main/java/com/facebook/presto/hive/metastore/thrift/HiveMetastoreClient.java +++ b/presto-hive-metastore/src/main/java/com/facebook/presto/hive/metastore/thrift/HiveMetastoreClient.java @@ -16,6 +16,7 @@ import org.apache.hadoop.hive.metastore.api.CheckLockRequest; import org.apache.hadoop.hive.metastore.api.ColumnStatisticsObj; import org.apache.hadoop.hive.metastore.api.Database; +import org.apache.hadoop.hive.metastore.api.EnvironmentContext; import org.apache.hadoop.hive.metastore.api.FieldSchema; import org.apache.hadoop.hive.metastore.api.HiveObjectPrivilege; import org.apache.hadoop.hive.metastore.api.HiveObjectRef; @@ -86,6 +87,9 @@ void dropTable(String databaseName, String name, boolean deleteData) void alterTable(String databaseName, String tableName, Table newTable) throws TException; + void alterTableWithEnvironmentContext(String databaseName, String tableName, Table newTable, EnvironmentContext context) + throws TException; + Table getTable(String databaseName, String tableName) throws TException; diff --git a/presto-hive-metastore/src/main/java/com/facebook/presto/hive/metastore/thrift/ThriftHiveMetastore.java b/presto-hive-metastore/src/main/java/com/facebook/presto/hive/metastore/thrift/ThriftHiveMetastore.java index 7df3f397e4cc2..c3bb696979e17 100644 --- a/presto-hive-metastore/src/main/java/com/facebook/presto/hive/metastore/thrift/ThriftHiveMetastore.java +++ b/presto-hive-metastore/src/main/java/com/facebook/presto/hive/metastore/thrift/ThriftHiveMetastore.java @@ -64,6 +64,7 @@ import org.apache.hadoop.hive.metastore.api.CheckLockRequest; import org.apache.hadoop.hive.metastore.api.ColumnStatisticsObj; import org.apache.hadoop.hive.metastore.api.Database; +import org.apache.hadoop.hive.metastore.api.EnvironmentContext; import org.apache.hadoop.hive.metastore.api.FieldSchema; import org.apache.hadoop.hive.metastore.api.HiveObjectPrivilege; import org.apache.hadoop.hive.metastore.api.HiveObjectRef; @@ -1160,6 +1161,35 @@ public MetastoreOperationResult alterTable(MetastoreContext metastoreContext, St } } + @Override + public MetastoreOperationResult alterTableWithEnvironmentContext(MetastoreContext metastoreContext, String databaseName, String tableName, Table table, EnvironmentContext environmentContext) + { + try { + retry() + .stopOn(InvalidOperationException.class, MetaException.class) + .stopOnIllegalExceptions() + .run("alterTableWithEnvironmentContext", stats.getAlterTableWithEnvironmentContext().wrap(() -> + getMetastoreClientThenCall(metastoreContext, client -> { + Optional source = getTable(metastoreContext, databaseName, tableName); + if (!source.isPresent()) { + throw new TableNotFoundException(new SchemaTableName(databaseName, tableName)); + } + client.alterTableWithEnvironmentContext(databaseName, tableName, table, environmentContext); + return null; + }))); + return EMPTY_RESULT; + } + catch (NoSuchObjectException e) { + throw new TableNotFoundException(new SchemaTableName(databaseName, tableName)); + } + catch (TException e) { + throw new PrestoException(HIVE_METASTORE_ERROR, e); + } + catch (Exception e) { + throw propagate(e); + } + } + @Override public Optional> getPartitionNames(MetastoreContext metastoreContext, String databaseName, String tableName) { diff --git a/presto-hive-metastore/src/main/java/com/facebook/presto/hive/metastore/thrift/ThriftHiveMetastoreClient.java b/presto-hive-metastore/src/main/java/com/facebook/presto/hive/metastore/thrift/ThriftHiveMetastoreClient.java index f34955d09432e..93ed690cc87a6 100644 --- a/presto-hive-metastore/src/main/java/com/facebook/presto/hive/metastore/thrift/ThriftHiveMetastoreClient.java +++ b/presto-hive-metastore/src/main/java/com/facebook/presto/hive/metastore/thrift/ThriftHiveMetastoreClient.java @@ -23,6 +23,7 @@ import org.apache.hadoop.hive.metastore.api.ColumnStatisticsObj; import org.apache.hadoop.hive.metastore.api.Database; import org.apache.hadoop.hive.metastore.api.DropConstraintRequest; +import org.apache.hadoop.hive.metastore.api.EnvironmentContext; import org.apache.hadoop.hive.metastore.api.FieldSchema; import org.apache.hadoop.hive.metastore.api.GetRoleGrantsForPrincipalRequest; import org.apache.hadoop.hive.metastore.api.GetRoleGrantsForPrincipalResponse; @@ -199,6 +200,13 @@ public void alterTable(String databaseName, String tableName, Table newTable) client.alter_table(constructSchemaName(catalogName, databaseName), tableName, newTable); } + @Override + public void alterTableWithEnvironmentContext(String databaseName, String tableName, Table newTable, EnvironmentContext context) + throws TException + { + client.alter_table_with_environment_context(constructSchemaName(catalogName, databaseName), tableName, newTable, context); + } + @Override public Table getTable(String databaseName, String tableName) throws TException diff --git a/presto-hive-metastore/src/main/java/com/facebook/presto/hive/metastore/thrift/ThriftHiveMetastoreStats.java b/presto-hive-metastore/src/main/java/com/facebook/presto/hive/metastore/thrift/ThriftHiveMetastoreStats.java index 543c1a2ba27da..aadc7f225b097 100644 --- a/presto-hive-metastore/src/main/java/com/facebook/presto/hive/metastore/thrift/ThriftHiveMetastoreStats.java +++ b/presto-hive-metastore/src/main/java/com/facebook/presto/hive/metastore/thrift/ThriftHiveMetastoreStats.java @@ -40,6 +40,7 @@ public class ThriftHiveMetastoreStats private final HiveMetastoreApiStats createTableWithConstraints = new HiveMetastoreApiStats(); private final HiveMetastoreApiStats dropTable = new HiveMetastoreApiStats(); private final HiveMetastoreApiStats alterTable = new HiveMetastoreApiStats(); + private final HiveMetastoreApiStats alterTableWithEnvironmentContext = new HiveMetastoreApiStats(); private final HiveMetastoreApiStats addPartitions = new HiveMetastoreApiStats(); private final HiveMetastoreApiStats dropPartition = new HiveMetastoreApiStats(); private final HiveMetastoreApiStats alterPartition = new HiveMetastoreApiStats(); @@ -216,6 +217,13 @@ public HiveMetastoreApiStats getAlterTable() return alterTable; } + @Managed + @Nested + public HiveMetastoreApiStats getAlterTableWithEnvironmentContext() + { + return alterTableWithEnvironmentContext; + } + @Managed @Nested public HiveMetastoreApiStats getAddPartitions() diff --git a/presto-hive-metastore/src/test/java/com/facebook/presto/hive/metastore/UnimplementedHiveMetastore.java b/presto-hive-metastore/src/test/java/com/facebook/presto/hive/metastore/UnimplementedHiveMetastore.java index 018ed253bdb0f..2e3bbe18952f5 100644 --- a/presto-hive-metastore/src/test/java/com/facebook/presto/hive/metastore/UnimplementedHiveMetastore.java +++ b/presto-hive-metastore/src/test/java/com/facebook/presto/hive/metastore/UnimplementedHiveMetastore.java @@ -28,6 +28,7 @@ import java.util.Optional; import java.util.Set; import java.util.function.Function; +import java.util.function.Supplier; public class UnimplementedHiveMetastore implements ExtendedHiveMetastore @@ -128,6 +129,12 @@ public MetastoreOperationResult replaceTable(MetastoreContext metastoreContext, throw new UnsupportedOperationException(); } + @Override + public MetastoreOperationResult persistTable(MetastoreContext metastoreContext, String databaseName, String tableName, Table newTable, PrincipalPrivileges principalPrivileges, Supplier update, Map additionalParameters) + { + throw new UnsupportedOperationException(); + } + @Override public MetastoreOperationResult renameTable(MetastoreContext metastoreContext, String databaseName, String tableName, String newDatabaseName, String newTableName) { diff --git a/presto-hive-metastore/src/test/java/com/facebook/presto/hive/metastore/thrift/InMemoryHiveMetastore.java b/presto-hive-metastore/src/test/java/com/facebook/presto/hive/metastore/thrift/InMemoryHiveMetastore.java index 636e99b38196e..866697dd088df 100644 --- a/presto-hive-metastore/src/test/java/com/facebook/presto/hive/metastore/thrift/InMemoryHiveMetastore.java +++ b/presto-hive-metastore/src/test/java/com/facebook/presto/hive/metastore/thrift/InMemoryHiveMetastore.java @@ -40,6 +40,7 @@ import org.apache.hadoop.fs.Path; import org.apache.hadoop.hive.metastore.TableType; import org.apache.hadoop.hive.metastore.api.Database; +import org.apache.hadoop.hive.metastore.api.EnvironmentContext; import org.apache.hadoop.hive.metastore.api.Partition; import org.apache.hadoop.hive.metastore.api.PrincipalPrivilegeSet; import org.apache.hadoop.hive.metastore.api.PrincipalType; @@ -283,6 +284,12 @@ public synchronized MetastoreOperationResult alterTable(MetastoreContext metasto return EMPTY_RESULT; } + @Override + public synchronized MetastoreOperationResult alterTableWithEnvironmentContext(MetastoreContext metastoreContext, String databaseName, String tableName, Table newTable, EnvironmentContext environmentContext) + { + return alterTable(metastoreContext, databaseName, tableName, newTable); + } + @Override public synchronized Optional> getAllTables(MetastoreContext metastoreContext, String databaseName) { diff --git a/presto-hive-metastore/src/test/java/com/facebook/presto/hive/metastore/thrift/MockHiveMetastoreClient.java b/presto-hive-metastore/src/test/java/com/facebook/presto/hive/metastore/thrift/MockHiveMetastoreClient.java index 4ee92ab0e2b47..ffb7a77670d65 100644 --- a/presto-hive-metastore/src/test/java/com/facebook/presto/hive/metastore/thrift/MockHiveMetastoreClient.java +++ b/presto-hive-metastore/src/test/java/com/facebook/presto/hive/metastore/thrift/MockHiveMetastoreClient.java @@ -29,6 +29,7 @@ import org.apache.hadoop.hive.metastore.api.CheckLockRequest; import org.apache.hadoop.hive.metastore.api.ColumnStatisticsObj; import org.apache.hadoop.hive.metastore.api.Database; +import org.apache.hadoop.hive.metastore.api.EnvironmentContext; import org.apache.hadoop.hive.metastore.api.FieldSchema; import org.apache.hadoop.hive.metastore.api.HiveObjectPrivilege; import org.apache.hadoop.hive.metastore.api.HiveObjectRef; @@ -365,6 +366,13 @@ public void alterTable(String databaseName, String tableName, Table newTable) throw new UnsupportedOperationException(); } + @Override + public void alterTableWithEnvironmentContext(String databaseName, String tableName, Table newTable, EnvironmentContext context) + throws TException + { + throw new UnsupportedOperationException(); + } + @Override public int addPartitions(List newPartitions) { diff --git a/presto-hive/pom.xml b/presto-hive/pom.xml index 0e6cdf23cb5b5..c5bc57b835548 100644 --- a/presto-hive/pom.xml +++ b/presto-hive/pom.xml @@ -485,7 +485,7 @@ org.objenesis objenesis - 2.6 + 3.4 test @@ -494,6 +494,12 @@ + + com.facebook.presto + presto-sql-invoked-functions-plugin + ${project.version} + test + diff --git a/presto-hive/src/main/java/com/facebook/presto/hive/HiveClientConfig.java b/presto-hive/src/main/java/com/facebook/presto/hive/HiveClientConfig.java index be0b6482d39f8..7f56998606165 100644 --- a/presto-hive/src/main/java/com/facebook/presto/hive/HiveClientConfig.java +++ b/presto-hive/src/main/java/com/facebook/presto/hive/HiveClientConfig.java @@ -107,7 +107,6 @@ public class HiveClientConfig private DataSize textMaxLineLength = new DataSize(100, MEGABYTE); private boolean assumeCanonicalPartitionKeys; - private boolean useOrcColumnNames; private double orcDefaultBloomFilterFpp = 0.05; private boolean rcfileOptimizedWriterEnabled = true; private boolean rcfileWriterValidate; @@ -720,19 +719,6 @@ public HiveClientConfig setS3FileSystemType(S3FileSystemType s3FileSystemType) return this; } - public boolean isUseOrcColumnNames() - { - return useOrcColumnNames; - } - - @Config("hive.orc.use-column-names") - @ConfigDescription("Access ORC columns using names from the file first, and fallback to Hive schema column names if not found to ensure backward compatibility with old data") - public HiveClientConfig setUseOrcColumnNames(boolean useOrcColumnNames) - { - this.useOrcColumnNames = useOrcColumnNames; - return this; - } - public double getOrcDefaultBloomFilterFpp() { return orcDefaultBloomFilterFpp; diff --git a/presto-hive/src/main/java/com/facebook/presto/hive/HiveMetadata.java b/presto-hive/src/main/java/com/facebook/presto/hive/HiveMetadata.java index e0d07e7ce4a58..d8940cb5e0b9c 100644 --- a/presto-hive/src/main/java/com/facebook/presto/hive/HiveMetadata.java +++ b/presto-hive/src/main/java/com/facebook/presto/hive/HiveMetadata.java @@ -109,7 +109,6 @@ import com.google.common.collect.ImmutableMap; import com.google.common.collect.ImmutableSet; import com.google.common.collect.ImmutableSortedMap; -import com.google.common.collect.Iterables; import com.google.common.collect.Lists; import com.google.common.collect.Maps; import com.google.common.collect.Sets; @@ -185,6 +184,7 @@ import static com.facebook.presto.hive.HiveColumnHandle.FILE_SIZE_COLUMN_NAME; import static com.facebook.presto.hive.HiveColumnHandle.PATH_COLUMN_NAME; import static com.facebook.presto.hive.HiveColumnHandle.ROW_ID_COLUMN_NAME; +import static com.facebook.presto.hive.HiveColumnHandle.rowIdColumnHandle; import static com.facebook.presto.hive.HiveColumnHandle.updateRowIdHandle; import static com.facebook.presto.hive.HiveErrorCode.HIVE_COLUMN_ORDER_MISMATCH; import static com.facebook.presto.hive.HiveErrorCode.HIVE_CONCURRENT_MODIFICATION_DETECTED; @@ -2648,8 +2648,8 @@ private List getOrComputePartitions(HiveTableLayoutHandle layoutH else { TupleDomain partitionColumnPredicate = layoutHandle.getPartitionColumnPredicate(); Predicate> predicate = convertToPredicate(partitionColumnPredicate); - List tableLayoutResults = getTableLayouts(session, tableHandle, new Constraint<>(partitionColumnPredicate, predicate), Optional.empty()); - return ((HiveTableLayoutHandle) Iterables.getOnlyElement(tableLayoutResults).getTableLayout().getHandle()).getPartitions().get(); + ConnectorTableLayoutResult tableLayoutResult = getTableLayoutForConstraint(session, tableHandle, new Constraint<>(partitionColumnPredicate, predicate), Optional.empty()); + return ((HiveTableLayoutHandle) tableLayoutResult.getTableLayout().getHandle()).getPartitions().get(); } } @@ -2726,7 +2726,11 @@ private String createTableLayoutString( } @Override - public List getTableLayouts(ConnectorSession session, ConnectorTableHandle tableHandle, Constraint constraint, Optional> desiredColumns) + public ConnectorTableLayoutResult getTableLayoutForConstraint( + ConnectorSession session, + ConnectorTableHandle tableHandle, + Constraint constraint, + Optional> desiredColumns) { HiveTableHandle handle = (HiveTableHandle) tableHandle; HivePartitionResult hivePartitionResult; @@ -2754,7 +2758,7 @@ public List getTableLayouts(ConnectorSession session String layoutString = createTableLayoutString(session, handle.getSchemaTableName(), hivePartitionResult.getBucketHandle(), hivePartitionResult.getBucketFilter(), TRUE_CONSTANT, domainPredicate); Optional> requestedColumns = desiredColumns.map(columns -> columns.stream().map(column -> (HiveColumnHandle) column).collect(toImmutableSet())); - return ImmutableList.of(new ConnectorTableLayoutResult( + return new ConnectorTableLayoutResult( getTableLayout( session, new HiveTableLayoutHandle.Builder() @@ -2777,7 +2781,7 @@ public List getTableLayouts(ConnectorSession session .setAppendRowNumberEnabled(false) .setHiveTableHandle(handle) .build()), - hivePartitionResult.getUnenforcedConstraint())); + hivePartitionResult.getUnenforcedConstraint()); } private static Subfield toSubfield(ColumnHandle columnHandle) @@ -2924,7 +2928,8 @@ && isOrderBasedExecutionEnabled(session)) { streamPartitionColumns, discretePredicates, localPropertyBuilder.build(), - Optional.of(combinedRemainingPredicate)); + Optional.of(combinedRemainingPredicate), + Optional.of(rowIdColumnHandle())); } @Override diff --git a/presto-hive/src/main/java/com/facebook/presto/hive/orc/DwrfAggregatedPageSourceFactory.java b/presto-hive/src/main/java/com/facebook/presto/hive/orc/DwrfAggregatedPageSourceFactory.java index 30fc1dcc49438..723759a0b162e 100644 --- a/presto-hive/src/main/java/com/facebook/presto/hive/orc/DwrfAggregatedPageSourceFactory.java +++ b/presto-hive/src/main/java/com/facebook/presto/hive/orc/DwrfAggregatedPageSourceFactory.java @@ -19,7 +19,6 @@ import com.facebook.presto.hive.FileFormatDataSourceStats; import com.facebook.presto.hive.HdfsEnvironment; import com.facebook.presto.hive.HiveAggregatedPageSourceFactory; -import com.facebook.presto.hive.HiveClientConfig; import com.facebook.presto.hive.HiveColumnHandle; import com.facebook.presto.hive.HiveFileContext; import com.facebook.presto.hive.HiveFileSplit; @@ -36,6 +35,7 @@ import java.util.List; import java.util.Optional; +import static com.facebook.presto.hive.HiveCommonSessionProperties.isUseOrcColumnNames; import static com.facebook.presto.hive.HiveErrorCode.HIVE_BAD_DATA; import static com.facebook.presto.hive.orc.OrcAggregatedPageSourceFactory.createOrcPageSource; import static com.facebook.presto.orc.DwrfEncryptionProvider.NO_ENCRYPTION; @@ -47,7 +47,6 @@ public class DwrfAggregatedPageSourceFactory { private final TypeManager typeManager; private final StandardFunctionResolution functionResolution; - private final boolean useOrcColumnNames; private final HdfsEnvironment hdfsEnvironment; private final FileFormatDataSourceStats stats; private final OrcFileTailSource orcFileTailSource; @@ -57,26 +56,6 @@ public class DwrfAggregatedPageSourceFactory public DwrfAggregatedPageSourceFactory( TypeManager typeManager, StandardFunctionResolution functionResolution, - HiveClientConfig config, - HdfsEnvironment hdfsEnvironment, - FileFormatDataSourceStats stats, - OrcFileTailSource orcFileTailSource, - StripeMetadataSourceFactory stripeMetadataSourceFactory) - { - this( - typeManager, - functionResolution, - requireNonNull(config, "hiveClientConfig is null").isUseOrcColumnNames(), - hdfsEnvironment, - stats, - orcFileTailSource, - stripeMetadataSourceFactory); - } - - public DwrfAggregatedPageSourceFactory( - TypeManager typeManager, - StandardFunctionResolution functionResolution, - boolean useOrcColumnNames, HdfsEnvironment hdfsEnvironment, FileFormatDataSourceStats stats, OrcFileTailSource orcFileTailSource, @@ -84,7 +63,6 @@ public DwrfAggregatedPageSourceFactory( { this.typeManager = requireNonNull(typeManager, "typeManager is null"); this.functionResolution = requireNonNull(functionResolution, "functionResolution is null"); - this.useOrcColumnNames = useOrcColumnNames; this.hdfsEnvironment = requireNonNull(hdfsEnvironment, "hdfsEnvironment is null"); this.stats = requireNonNull(stats, "stats is null"); this.orcFileTailSource = requireNonNull(orcFileTailSource, "orcFileTailCache is null"); @@ -116,7 +94,7 @@ public Optional createPageSource( configuration, fileSplit, columns, - useOrcColumnNames, + isUseOrcColumnNames(session), typeManager, functionResolution, stats, diff --git a/presto-hive/src/main/java/com/facebook/presto/hive/orc/OrcAggregatedPageSourceFactory.java b/presto-hive/src/main/java/com/facebook/presto/hive/orc/OrcAggregatedPageSourceFactory.java index c61507c7983bb..cd6313d484755 100644 --- a/presto-hive/src/main/java/com/facebook/presto/hive/orc/OrcAggregatedPageSourceFactory.java +++ b/presto-hive/src/main/java/com/facebook/presto/hive/orc/OrcAggregatedPageSourceFactory.java @@ -19,7 +19,6 @@ import com.facebook.presto.hive.FileFormatDataSourceStats; import com.facebook.presto.hive.HdfsEnvironment; import com.facebook.presto.hive.HiveAggregatedPageSourceFactory; -import com.facebook.presto.hive.HiveClientConfig; import com.facebook.presto.hive.HiveColumnHandle; import com.facebook.presto.hive.HiveFileContext; import com.facebook.presto.hive.HiveFileSplit; @@ -49,6 +48,7 @@ import static com.facebook.presto.hive.HiveCommonSessionProperties.getOrcMaxReadBlockSize; import static com.facebook.presto.hive.HiveCommonSessionProperties.getOrcTinyStripeThreshold; import static com.facebook.presto.hive.HiveCommonSessionProperties.isOrcZstdJniDecompressionEnabled; +import static com.facebook.presto.hive.HiveCommonSessionProperties.isUseOrcColumnNames; import static com.facebook.presto.hive.HiveUtil.getPhysicalHiveColumnHandles; import static com.facebook.presto.hive.orc.OrcPageSourceFactoryUtils.getOrcDataSource; import static com.facebook.presto.hive.orc.OrcPageSourceFactoryUtils.getOrcReader; @@ -62,7 +62,6 @@ public class OrcAggregatedPageSourceFactory { private final TypeManager typeManager; private final StandardFunctionResolution functionResolution; - private final boolean useOrcColumnNames; private final HdfsEnvironment hdfsEnvironment; private final FileFormatDataSourceStats stats; private final OrcFileTailSource orcFileTailSource; @@ -72,26 +71,6 @@ public class OrcAggregatedPageSourceFactory public OrcAggregatedPageSourceFactory( TypeManager typeManager, StandardFunctionResolution functionResolution, - HiveClientConfig config, - HdfsEnvironment hdfsEnvironment, - FileFormatDataSourceStats stats, - OrcFileTailSource orcFileTailSource, - StripeMetadataSourceFactory stripeMetadataSourceFactory) - { - this( - typeManager, - functionResolution, - requireNonNull(config, "hiveClientConfig is null").isUseOrcColumnNames(), - hdfsEnvironment, - stats, - orcFileTailSource, - stripeMetadataSourceFactory); - } - - public OrcAggregatedPageSourceFactory( - TypeManager typeManager, - StandardFunctionResolution functionResolution, - boolean useOrcColumnNames, HdfsEnvironment hdfsEnvironment, FileFormatDataSourceStats stats, OrcFileTailSource orcFileTailSource, @@ -99,7 +78,6 @@ public OrcAggregatedPageSourceFactory( { this.typeManager = requireNonNull(typeManager, "typeManager is null"); this.functionResolution = requireNonNull(functionResolution, "functionResolution is null"); - this.useOrcColumnNames = useOrcColumnNames; this.hdfsEnvironment = requireNonNull(hdfsEnvironment, "hdfsEnvironment is null"); this.stats = requireNonNull(stats, "stats is null"); this.orcFileTailSource = requireNonNull(orcFileTailSource, "orcFileTailCache is null"); @@ -132,7 +110,7 @@ public Optional createPageSource( configuration, fileSplit, columns, - useOrcColumnNames, + isUseOrcColumnNames(session), typeManager, functionResolution, stats, diff --git a/presto-hive/src/main/java/com/facebook/presto/hive/orc/OrcBatchPageSourceFactory.java b/presto-hive/src/main/java/com/facebook/presto/hive/orc/OrcBatchPageSourceFactory.java index 3c92bf8efc405..f1b892d4495b2 100644 --- a/presto-hive/src/main/java/com/facebook/presto/hive/orc/OrcBatchPageSourceFactory.java +++ b/presto-hive/src/main/java/com/facebook/presto/hive/orc/OrcBatchPageSourceFactory.java @@ -62,6 +62,7 @@ import static com.facebook.presto.hive.HiveCommonSessionProperties.getOrcTinyStripeThreshold; import static com.facebook.presto.hive.HiveCommonSessionProperties.isOrcBloomFiltersEnabled; import static com.facebook.presto.hive.HiveCommonSessionProperties.isOrcZstdJniDecompressionEnabled; +import static com.facebook.presto.hive.HiveCommonSessionProperties.isUseOrcColumnNames; import static com.facebook.presto.hive.HiveUtil.checkRowIDPartitionComponent; import static com.facebook.presto.hive.HiveUtil.getPhysicalHiveColumnHandles; import static com.facebook.presto.hive.orc.OrcPageSourceFactoryUtils.getOrcDataSource; @@ -78,7 +79,6 @@ public class OrcBatchPageSourceFactory implements HiveBatchPageSourceFactory { private final TypeManager typeManager; - private final boolean useOrcColumnNames; private final HdfsEnvironment hdfsEnvironment; private final FileFormatDataSourceStats stats; private final int domainCompactionThreshold; @@ -97,7 +97,6 @@ public OrcBatchPageSourceFactory( { this( typeManager, - requireNonNull(config, "hiveClientConfig is null").isUseOrcColumnNames(), hdfsEnvironment, stats, config.getDomainCompactionThreshold(), @@ -107,7 +106,6 @@ public OrcBatchPageSourceFactory( public OrcBatchPageSourceFactory( TypeManager typeManager, - boolean useOrcColumnNames, HdfsEnvironment hdfsEnvironment, FileFormatDataSourceStats stats, int domainCompactionThreshold, @@ -115,7 +113,6 @@ public OrcBatchPageSourceFactory( StripeMetadataSourceFactory stripeMetadataSourceFactory) { this.typeManager = requireNonNull(typeManager, "typeManager is null"); - this.useOrcColumnNames = useOrcColumnNames; this.hdfsEnvironment = requireNonNull(hdfsEnvironment, "hdfsEnvironment is null"); this.stats = requireNonNull(stats, "stats is null"); this.domainCompactionThreshold = domainCompactionThreshold; @@ -153,7 +150,7 @@ public Optional createPageSource( configuration, fileSplit, columns, - useOrcColumnNames, + isUseOrcColumnNames(session), effectivePredicate, hiveStorageTimeZone, typeManager, diff --git a/presto-hive/src/main/java/com/facebook/presto/hive/orc/OrcSelectivePageSourceFactory.java b/presto-hive/src/main/java/com/facebook/presto/hive/orc/OrcSelectivePageSourceFactory.java index 60354406ace2f..5fcdc9c89ad7c 100644 --- a/presto-hive/src/main/java/com/facebook/presto/hive/orc/OrcSelectivePageSourceFactory.java +++ b/presto-hive/src/main/java/com/facebook/presto/hive/orc/OrcSelectivePageSourceFactory.java @@ -104,6 +104,7 @@ import static com.facebook.presto.hive.HiveCommonSessionProperties.getOrcTinyStripeThreshold; import static com.facebook.presto.hive.HiveCommonSessionProperties.isOrcBloomFiltersEnabled; import static com.facebook.presto.hive.HiveCommonSessionProperties.isOrcZstdJniDecompressionEnabled; +import static com.facebook.presto.hive.HiveCommonSessionProperties.isUseOrcColumnNames; import static com.facebook.presto.hive.HiveErrorCode.HIVE_INVALID_BUCKET_FILES; import static com.facebook.presto.hive.HiveSessionProperties.isAdaptiveFilterReorderingEnabled; import static com.facebook.presto.hive.HiveSessionProperties.isLegacyTimestampBucketing; @@ -132,7 +133,6 @@ public class OrcSelectivePageSourceFactory private final TypeManager typeManager; private final StandardFunctionResolution functionResolution; private final RowExpressionService rowExpressionService; - private final boolean useOrcColumnNames; private final HdfsEnvironment hdfsEnvironment; private final FileFormatDataSourceStats stats; private final int domainCompactionThreshold; @@ -156,7 +156,6 @@ public OrcSelectivePageSourceFactory( typeManager, functionResolution, rowExpressionService, - requireNonNull(config, "hiveClientConfig is null").isUseOrcColumnNames(), hdfsEnvironment, stats, config.getDomainCompactionThreshold(), @@ -169,7 +168,6 @@ public OrcSelectivePageSourceFactory( TypeManager typeManager, StandardFunctionResolution functionResolution, RowExpressionService rowExpressionService, - boolean useOrcColumnNames, HdfsEnvironment hdfsEnvironment, FileFormatDataSourceStats stats, int domainCompactionThreshold, @@ -180,7 +178,6 @@ public OrcSelectivePageSourceFactory( this.typeManager = requireNonNull(typeManager, "typeManager is null"); this.functionResolution = requireNonNull(functionResolution, "functionResolution is null"); this.rowExpressionService = requireNonNull(rowExpressionService, "rowExpressionService is null"); - this.useOrcColumnNames = useOrcColumnNames; this.hdfsEnvironment = requireNonNull(hdfsEnvironment, "hdfsEnvironment is null"); this.stats = requireNonNull(stats, "stats is null"); this.domainCompactionThreshold = domainCompactionThreshold; @@ -230,7 +227,7 @@ public Optional createPageSource( outputColumns, domainPredicate, remainingPredicate, - useOrcColumnNames, + isUseOrcColumnNames(session), hiveStorageTimeZone, typeManager, functionResolution, diff --git a/presto-hive/src/test/java/com/facebook/presto/hive/AbstractTestHiveClient.java b/presto-hive/src/test/java/com/facebook/presto/hive/AbstractTestHiveClient.java index 17b2e70cff0ec..38bd3ece02ad3 100644 --- a/presto-hive/src/test/java/com/facebook/presto/hive/AbstractTestHiveClient.java +++ b/presto-hive/src/test/java/com/facebook/presto/hive/AbstractTestHiveClient.java @@ -2328,7 +2328,7 @@ private void doTestBucketedTableEvolution(HiveStorageFormat storageFormat, Schem Optional.empty()).getLayout().getHandle(); } else { - layoutHandle = getOnlyElement(metadata.getTableLayouts(session, tableHandle, new Constraint<>(TupleDomain.fromFixedValues(ImmutableMap.of(bucketColumnHandle(), singleBucket))), Optional.empty())).getTableLayout().getHandle(); + layoutHandle = metadata.getTableLayoutForConstraint(session, tableHandle, new Constraint<>(TupleDomain.fromFixedValues(ImmutableMap.of(bucketColumnHandle(), singleBucket))), Optional.empty()).getTableLayout().getHandle(); } result = readTable( @@ -2686,8 +2686,8 @@ protected ConnectorTableLayout getTableLayout(ConnectorSession session, Connecto Optional.empty()).getLayout(); } - List tableLayoutResults = metadata.getTableLayouts(session, tableHandle, constraint, Optional.empty()); - return getOnlyElement(tableLayoutResults).getTableLayout(); + ConnectorTableLayoutResult tableLayoutResult = metadata.getTableLayoutForConstraint(session, tableHandle, constraint, Optional.empty()); + return tableLayoutResult.getTableLayout(); } @Test diff --git a/presto-hive/src/test/java/com/facebook/presto/hive/AbstractTestHiveFileSystem.java b/presto-hive/src/test/java/com/facebook/presto/hive/AbstractTestHiveFileSystem.java index 5ec04cc1aa4c3..866d3f797d7fc 100644 --- a/presto-hive/src/test/java/com/facebook/presto/hive/AbstractTestHiveFileSystem.java +++ b/presto-hive/src/test/java/com/facebook/presto/hive/AbstractTestHiveFileSystem.java @@ -470,8 +470,8 @@ private void createTable(MetastoreContext metastoreContext, SchemaTableName tabl assertEquals(filterNonHiddenColumnMetadata(tableMetadata.getColumns()), columns); // verify the data - List tableLayoutResults = metadata.getTableLayouts(session, hiveTableHandle, Constraint.alwaysTrue(), Optional.empty()); - HiveTableLayoutHandle layoutHandle = (HiveTableLayoutHandle) getOnlyElement(tableLayoutResults).getTableLayout().getHandle(); + ConnectorTableLayoutResult tableLayoutResult = metadata.getTableLayoutForConstraint(session, hiveTableHandle, Constraint.alwaysTrue(), Optional.empty()); + HiveTableLayoutHandle layoutHandle = (HiveTableLayoutHandle) tableLayoutResult.getTableLayout().getHandle(); assertEquals(layoutHandle.getPartitions().get().size(), 1); ConnectorSplitSource splitSource = splitManager.getSplits(transaction.getTransactionHandle(), session, layoutHandle, SPLIT_SCHEDULING_CONTEXT); ConnectorSplit split = getOnlyElement(getAllSplits(splitSource)); diff --git a/presto-hive/src/test/java/com/facebook/presto/hive/HiveFileSystemTestUtils.java b/presto-hive/src/test/java/com/facebook/presto/hive/HiveFileSystemTestUtils.java index 0d509285ed16a..41f5000b446be 100644 --- a/presto-hive/src/test/java/com/facebook/presto/hive/HiveFileSystemTestUtils.java +++ b/presto-hive/src/test/java/com/facebook/presto/hive/HiveFileSystemTestUtils.java @@ -50,7 +50,6 @@ import static com.facebook.presto.testing.MaterializedResult.materializeSourceDataStream; import static com.google.common.base.Preconditions.checkArgument; import static com.google.common.collect.ImmutableList.toImmutableList; -import static com.google.common.collect.Iterables.getOnlyElement; public class HiveFileSystemTestUtils { @@ -74,8 +73,8 @@ public static MaterializedResult readTable(SchemaTableName tableName, ConnectorTableHandle table = getTableHandle(metadata, tableName, session); List columnHandles = ImmutableList.copyOf(metadata.getColumnHandles(session, table).values()); - List tableLayoutResults = metadata.getTableLayouts(session, table, Constraint.alwaysTrue(), Optional.empty()); - HiveTableLayoutHandle layoutHandle = (HiveTableLayoutHandle) getOnlyElement(tableLayoutResults).getTableLayout().getHandle(); + ConnectorTableLayoutResult tableLayoutResult = metadata.getTableLayoutForConstraint(session, table, Constraint.alwaysTrue(), Optional.empty()); + HiveTableLayoutHandle layoutHandle = (HiveTableLayoutHandle) tableLayoutResult.getTableLayout().getHandle(); TableHandle tableHandle = new TableHandle(new ConnectorId(tableName.getSchemaName()), table, transaction.getTransactionHandle(), Optional.of(layoutHandle)); metadata.beginQuery(session); @@ -134,8 +133,8 @@ public static MaterializedResult filterTable(SchemaTableName tableName, session = newSession(config); ConnectorTableHandle table = getTableHandle(metadata, tableName, session); - List tableLayoutResults = metadata.getTableLayouts(session, table, Constraint.alwaysTrue(), Optional.empty()); - HiveTableLayoutHandle layoutHandle = (HiveTableLayoutHandle) getOnlyElement(tableLayoutResults).getTableLayout().getHandle(); + ConnectorTableLayoutResult tableLayoutResult = metadata.getTableLayoutForConstraint(session, table, Constraint.alwaysTrue(), Optional.empty()); + HiveTableLayoutHandle layoutHandle = (HiveTableLayoutHandle) tableLayoutResult.getTableLayout().getHandle(); TableHandle tableHandle = new TableHandle(new ConnectorId(tableName.getSchemaName()), table, transaction.getTransactionHandle(), Optional.of(layoutHandle)); metadata.beginQuery(session); @@ -190,8 +189,8 @@ public static int getSplitsCount(SchemaTableName tableName, session = newSession(config); ConnectorTableHandle table = getTableHandle(metadata, tableName, session); - List tableLayoutResults = metadata.getTableLayouts(session, table, Constraint.alwaysTrue(), Optional.empty()); - HiveTableLayoutHandle layoutHandle = (HiveTableLayoutHandle) getOnlyElement(tableLayoutResults).getTableLayout().getHandle(); + ConnectorTableLayoutResult tableLayoutResult = metadata.getTableLayoutForConstraint(session, table, Constraint.alwaysTrue(), Optional.empty()); + HiveTableLayoutHandle layoutHandle = (HiveTableLayoutHandle) tableLayoutResult.getTableLayout().getHandle(); TableHandle tableHandle = new TableHandle(new ConnectorId(tableName.getSchemaName()), table, transaction.getTransactionHandle(), Optional.of(layoutHandle)); metadata.beginQuery(session); diff --git a/presto-hive/src/test/java/com/facebook/presto/hive/HiveTestUtils.java b/presto-hive/src/test/java/com/facebook/presto/hive/HiveTestUtils.java index 37b3b35139a27..14d00eee6dc81 100644 --- a/presto-hive/src/test/java/com/facebook/presto/hive/HiveTestUtils.java +++ b/presto-hive/src/test/java/com/facebook/presto/hive/HiveTestUtils.java @@ -17,7 +17,9 @@ import com.facebook.airlift.json.smile.SmileCodec; import com.facebook.airlift.log.Logger; import com.facebook.presto.PagesIndexPageSorter; +import com.facebook.presto.Session; import com.facebook.presto.cache.CacheConfig; +import com.facebook.presto.common.QualifiedObjectName; import com.facebook.presto.common.block.BlockEncodingManager; import com.facebook.presto.common.type.ArrayType; import com.facebook.presto.common.type.MapType; @@ -51,7 +53,9 @@ import com.facebook.presto.hive.s3.PrestoS3ConfigurationUpdater; import com.facebook.presto.hive.s3select.S3SelectRecordCursorProvider; import com.facebook.presto.metadata.FunctionAndTypeManager; +import com.facebook.presto.metadata.Metadata; import com.facebook.presto.metadata.MetadataManager; +import com.facebook.presto.metadata.TableLayout; import com.facebook.presto.operator.PagesIndex; import com.facebook.presto.orc.StorageStripeMetadataSource; import com.facebook.presto.orc.StripeMetadataSourceFactory; @@ -59,7 +63,9 @@ import com.facebook.presto.parquet.cache.MetadataReader; import com.facebook.presto.spi.ColumnHandle; import com.facebook.presto.spi.ConnectorSession; +import com.facebook.presto.spi.Constraint; import com.facebook.presto.spi.PageSorter; +import com.facebook.presto.spi.TableHandle; import com.facebook.presto.spi.function.StandardFunctionResolution; import com.facebook.presto.spi.plan.FilterStatsCalculatorService; import com.facebook.presto.spi.relation.DeterminismEvaluator; @@ -75,7 +81,9 @@ import com.facebook.presto.sql.relational.RowExpressionDeterminismEvaluator; import com.facebook.presto.sql.relational.RowExpressionDomainTranslator; import com.facebook.presto.sql.relational.RowExpressionOptimizer; +import com.facebook.presto.testing.QueryRunner; import com.facebook.presto.testing.TestingConnectorSession; +import com.facebook.presto.tests.DistributedQueryRunner; import com.google.common.collect.ImmutableList; import com.google.common.collect.ImmutableSet; import io.airlift.slice.Slice; @@ -87,13 +95,17 @@ import java.util.List; import java.util.Optional; import java.util.Set; +import java.util.function.Function; import static com.facebook.airlift.json.JsonCodec.jsonCodec; import static com.facebook.airlift.json.smile.SmileCodec.smileCodec; import static com.facebook.presto.common.type.Decimals.encodeScaledValue; import static com.facebook.presto.hive.HiveDwrfEncryptionProvider.NO_ENCRYPTION; +import static com.facebook.presto.hive.HiveQueryRunner.TPCH_SCHEMA; +import static com.facebook.presto.transaction.TransactionBuilder.transaction; import static java.lang.String.format; import static java.util.stream.Collectors.toList; +import static org.testng.Assert.assertTrue; public final class HiveTestUtils { @@ -188,8 +200,8 @@ public static Set getDefaultHiveAggregatedPageS FileFormatDataSourceStats stats = new FileFormatDataSourceStats(); HdfsEnvironment testHdfsEnvironment = createTestHdfsEnvironment(hiveClientConfig, metastoreClientConfig); return ImmutableSet.builder() - .add(new OrcAggregatedPageSourceFactory(FUNCTION_AND_TYPE_MANAGER, FUNCTION_RESOLUTION, hiveClientConfig, testHdfsEnvironment, stats, new StorageOrcFileTailSource(), StripeMetadataSourceFactory.of(new StorageStripeMetadataSource()))) - .add(new DwrfAggregatedPageSourceFactory(FUNCTION_AND_TYPE_MANAGER, FUNCTION_RESOLUTION, hiveClientConfig, testHdfsEnvironment, stats, new StorageOrcFileTailSource(), StripeMetadataSourceFactory.of(new StorageStripeMetadataSource()))) + .add(new OrcAggregatedPageSourceFactory(FUNCTION_AND_TYPE_MANAGER, FUNCTION_RESOLUTION, testHdfsEnvironment, stats, new StorageOrcFileTailSource(), StripeMetadataSourceFactory.of(new StorageStripeMetadataSource()))) + .add(new DwrfAggregatedPageSourceFactory(FUNCTION_AND_TYPE_MANAGER, FUNCTION_RESOLUTION, testHdfsEnvironment, stats, new StorageOrcFileTailSource(), StripeMetadataSourceFactory.of(new StorageStripeMetadataSource()))) .add(new ParquetAggregatedPageSourceFactory(FUNCTION_AND_TYPE_MANAGER, FUNCTION_RESOLUTION, testHdfsEnvironment, stats, new MetadataReader())) .build(); } @@ -374,4 +386,20 @@ public static List> getAllSessionProperties(HiveClientConfig allSessionProperties.addAll(hiveCommonSessionProperties.getSessionProperties()); return allSessionProperties; } + + public static Object getHiveTableProperty(QueryRunner queryRunner, Session session, String tableName, Function propertyGetter) + { + Metadata metadata = ((DistributedQueryRunner) queryRunner).getCoordinator().getMetadata(); + + return transaction(queryRunner.getTransactionManager(), queryRunner.getAccessControl()) + .readOnly() + .execute(session, transactionSession -> { + Optional tableHandle = metadata.getMetadataResolver(transactionSession).getTableHandle(new QualifiedObjectName("hive", TPCH_SCHEMA, tableName)); + assertTrue(tableHandle.isPresent()); + + TableLayout layout = metadata.getLayout(transactionSession, tableHandle.get(), Constraint.alwaysTrue(), Optional.empty()) + .getLayout(); + return propertyGetter.apply((HiveTableLayoutHandle) layout.getNewTableHandle().getLayout().get()); + }); + } } diff --git a/presto-hive/src/test/java/com/facebook/presto/hive/TestDistributedQueriesSingleNode.java b/presto-hive/src/test/java/com/facebook/presto/hive/TestDistributedQueriesSingleNode.java index bfaec5f22852d..e3328a940c71b 100644 --- a/presto-hive/src/test/java/com/facebook/presto/hive/TestDistributedQueriesSingleNode.java +++ b/presto-hive/src/test/java/com/facebook/presto/hive/TestDistributedQueriesSingleNode.java @@ -13,6 +13,7 @@ */ package com.facebook.presto.hive; +import com.facebook.presto.scalar.sql.SqlInvokedFunctionsPlugin; import com.facebook.presto.testing.QueryRunner; import com.facebook.presto.tests.AbstractTestDistributedQueries; import com.google.common.collect.ImmutableMap; @@ -32,11 +33,13 @@ protected QueryRunner createQueryRunner() { ImmutableMap.Builder coordinatorProperties = ImmutableMap.builder(); coordinatorProperties.put("single-node-execution-enabled", "true"); - return HiveQueryRunner.createQueryRunner( + QueryRunner queryRunner = HiveQueryRunner.createQueryRunner( getTables(), ImmutableMap.of(), coordinatorProperties.build(), Optional.empty()); + queryRunner.installPlugin(new SqlInvokedFunctionsPlugin()); + return queryRunner; } @Override diff --git a/presto-hive/src/test/java/com/facebook/presto/hive/TestHiveBucketedTablesWithRowId.java b/presto-hive/src/test/java/com/facebook/presto/hive/TestHiveBucketedTablesWithRowId.java new file mode 100644 index 0000000000000..c32fa22ac1116 --- /dev/null +++ b/presto-hive/src/test/java/com/facebook/presto/hive/TestHiveBucketedTablesWithRowId.java @@ -0,0 +1,240 @@ +/* + * Licensed 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 com.facebook.presto.hive; + +import com.facebook.presto.Session; +import com.facebook.presto.testing.QueryRunner; +import com.facebook.presto.tests.AbstractTestQueryFramework; +import com.google.common.collect.ImmutableList; +import com.google.common.collect.ImmutableMap; +import org.testng.annotations.AfterClass; +import org.testng.annotations.BeforeClass; +import org.testng.annotations.Test; + +import java.util.Optional; + +import static com.facebook.presto.SystemSessionProperties.UTILIZE_UNIQUE_PROPERTY_IN_QUERY_PLANNING; +import static com.facebook.presto.sql.planner.assertions.PlanMatchPattern.aggregation; +import static com.facebook.presto.sql.planner.assertions.PlanMatchPattern.anyTree; +import static com.facebook.presto.sql.planner.assertions.PlanMatchPattern.exchange; +import static com.facebook.presto.sql.planner.assertions.PlanMatchPattern.filter; +import static com.facebook.presto.sql.planner.assertions.PlanMatchPattern.join; +import static com.facebook.presto.sql.planner.assertions.PlanMatchPattern.project; +import static com.facebook.presto.sql.planner.assertions.PlanMatchPattern.tableScan; +import static io.airlift.tpch.TpchTable.CUSTOMER; +import static io.airlift.tpch.TpchTable.ORDERS; + +@Test(singleThreaded = true) +public class TestHiveBucketedTablesWithRowId + extends AbstractTestQueryFramework +{ + @Override + protected QueryRunner createQueryRunner() + throws Exception + { + return HiveQueryRunner.createQueryRunner( + ImmutableList.of(ORDERS, CUSTOMER), + ImmutableMap.of(), + Optional.empty()); + } + + @BeforeClass + public void setUp() + { + // Create bucketed customer table + assertUpdate("CREATE TABLE customer_bucketed WITH " + + "(bucketed_by = ARRAY['custkey'], bucket_count = 13) " + + "AS SELECT * FROM customer", 1500); + + // Create bucketed orders table + assertUpdate("CREATE TABLE orders_bucketed WITH " + + "(bucketed_by = ARRAY['orderkey'], bucket_count = 11) " + + "AS SELECT * FROM orders", 15000); + + // Verify tables are created + assertQuery("SELECT count(*) FROM customer_bucketed", "SELECT count(*) FROM customer"); + assertQuery("SELECT count(*) FROM orders_bucketed", "SELECT count(*) FROM orders"); + } + + @AfterClass(alwaysRun = true) + public void tearDown() + { + try { + assertUpdate("DROP TABLE IF EXISTS customer_bucketed"); + assertUpdate("DROP TABLE IF EXISTS orders_bucketed"); + } + catch (Exception e) { + // Ignore cleanup errors + } + } + + @Test + public void testRowIdWithBucketColumn() + { + Session session = Session.builder(getQueryRunner().getDefaultSession()) + .setSystemProperty(UTILIZE_UNIQUE_PROPERTY_IN_QUERY_PLANNING, "true") + .build(); + + // Test basic query with both $row_id and $bucket + String sql = "SELECT \"$row_id\", \"$bucket\", custkey, name " + + "FROM customer_bucketed " + + "WHERE \"$bucket\" = 5"; + + assertPlan(session, sql, anyTree( + project(filter(tableScan("customer_bucketed"))))); + + // Test aggregation grouping by both $row_id and $bucket + sql = "SELECT \"$row_id\", \"$bucket\", COUNT(*) " + + "FROM customer_bucketed " + + "GROUP BY \"$row_id\", \"$bucket\""; + + assertPlan(session, sql, anyTree( + aggregation(ImmutableMap.of(), + project(tableScan("customer_bucketed"))))); + + // Test join between bucketed tables using both $row_id and $bucket + sql = "SELECT c.\"$row_id\" AS customer_row_id, " + + "c.\"$bucket\" AS customer_bucket, " + + "o.\"$row_id\" AS order_row_id, " + + "o.\"$bucket\" AS order_bucket, " + + "c.name, o.orderkey " + + "FROM customer_bucketed c " + + "JOIN orders_bucketed o " + + "ON c.custkey = o.custkey " + + "WHERE c.\"$bucket\" IN (1, 3, 5) " + + "AND o.\"$bucket\" IN (2, 4, 6)"; + + assertPlan(session, sql, anyTree( + join( + project(filter(tableScan("customer_bucketed"))), + exchange(anyTree(tableScan("orders_bucketed")))))); + } + + @Test + public void testRowIdUniquePropertyWithBucketing() + { + Session session = Session.builder(getQueryRunner().getDefaultSession()) + .setSystemProperty(UTILIZE_UNIQUE_PROPERTY_IN_QUERY_PLANNING, "true") + .build(); + + // Test unique grouping by $row_id with bucket filtering + String sql = "SELECT " + + "customer_row_id, " + + "ARBITRARY(name) AS customer_name, " + + "ARBITRARY(bucket_num) AS customer_bucket, " + + "ARRAY_AGG(orderkey) AS orders_info " + + "FROM (" + + " SELECT " + + " c.\"$row_id\" AS customer_row_id, " + + " c.\"$bucket\" AS bucket_num, " + + " c.name, " + + " o.orderkey " + + " FROM customer_bucketed c " + + " LEFT JOIN orders_bucketed o " + + " ON c.custkey = o.custkey " + + " AND o.orderstatus IN ('O', 'F') " + + " WHERE c.\"$bucket\" < 5 " + + " AND c.nationkey IN (1, 2, 3) " + + ") " + + "GROUP BY customer_row_id"; + + assertPlan(session, sql, anyTree( + aggregation(ImmutableMap.of(), + join( + anyTree(tableScan("customer_bucketed")), + anyTree(tableScan("orders_bucketed")))))); + } + + @Test + public void testRowIdAndBucketInComplexQuery() + { + Session session = Session.builder(getQueryRunner().getDefaultSession()) + .setSystemProperty(UTILIZE_UNIQUE_PROPERTY_IN_QUERY_PLANNING, "true") + .build(); + + // Complex query with both row_id and bucket columns + String sql = "SELECT " + + " unique_id, " + + " bucket_group, " + + " COUNT(*) AS order_count, " + + " AVG(totalprice) AS avg_price " + + "FROM (" + + " SELECT " + + " c.\"$row_id\" AS unique_id, " + + " CASE " + + " WHEN c.\"$bucket\" < 5 THEN 'low' " + + " WHEN c.\"$bucket\" < 10 THEN 'medium' " + + " ELSE 'high' " + + " END AS bucket_group, " + + " o.totalprice " + + " FROM customer_bucketed c " + + " JOIN orders_bucketed o " + + " ON c.custkey = o.custkey " + + " WHERE o.\"$bucket\" % 2 = 0 " + + ") t " + + "GROUP BY unique_id, bucket_group"; + + assertPlan(session, sql, anyTree( + aggregation(ImmutableMap.of(), + project(project(join( + project(tableScan("customer_bucketed")), + anyTree(tableScan("orders_bucketed")))))))); + } + + @Test + public void testDistinctRowIdWithBucketFilter() + { + Session session = Session.builder(getQueryRunner().getDefaultSession()) + .setSystemProperty(UTILIZE_UNIQUE_PROPERTY_IN_QUERY_PLANNING, "true") + .build(); + + // Test DISTINCT with row_id and bucket filtering + String sql = "SELECT " + + " DISTINCT c.\"$row_id\" AS unique_id, " + + " c.\"$bucket\" AS bucket_num, " + + " c.name " + + "FROM customer_bucketed c " + + "WHERE c.\"$bucket\" BETWEEN 3 AND 8 " + + " AND c.nationkey = 1"; + + assertPlan(session, sql, anyTree( + aggregation(ImmutableMap.of(), + project(filter(tableScan("customer_bucketed")))))); + } + + @Test + public void testRowIdJoinOnBucketColumn() + { + Session session = Session.builder(getQueryRunner().getDefaultSession()) + .setSystemProperty(UTILIZE_UNIQUE_PROPERTY_IN_QUERY_PLANNING, "true") + .build(); + + // Test joining on bucket column while selecting row_id + String sql = "SELECT " + + " c.\"$row_id\" AS customer_row_id, " + + " o.\"$row_id\" AS order_row_id, " + + " c.\"$bucket\" AS shared_bucket, " + + " c.name, " + + " o.orderkey " + + "FROM customer_bucketed c " + + "JOIN orders_bucketed o " + + " ON c.\"$bucket\" = o.\"$bucket\" " + + "WHERE c.\"$bucket\" < 5"; + + assertPlan(session, sql, anyTree( + join( + exchange(anyTree(tableScan("customer_bucketed"))), + exchange(anyTree(tableScan("orders_bucketed")))))); + } +} diff --git a/presto-hive/src/test/java/com/facebook/presto/hive/TestHiveClientConfig.java b/presto-hive/src/test/java/com/facebook/presto/hive/TestHiveClientConfig.java index 878330c46764e..1d09a4993dff6 100644 --- a/presto-hive/src/test/java/com/facebook/presto/hive/TestHiveClientConfig.java +++ b/presto-hive/src/test/java/com/facebook/presto/hive/TestHiveClientConfig.java @@ -83,7 +83,6 @@ public void testDefaults() .setMaxPartitionsPerWriter(100) .setWriteValidationThreads(16) .setTextMaxLineLength(new DataSize(100, Unit.MEGABYTE)) - .setUseOrcColumnNames(false) .setAssumeCanonicalPartitionKeys(false) .setOrcDefaultBloomFilterFpp(0.05) .setRcfileOptimizedWriterEnabled(true) @@ -211,7 +210,6 @@ public void testExplicitPropertyMappings() .put("hive.max-concurrent-zero-row-file-creations", "100") .put("hive.assume-canonical-partition-keys", "true") .put("hive.text.max-line-length", "13MB") - .put("hive.orc.use-column-names", "true") .put("hive.orc.default-bloom-filter-fpp", "0.96") .put("hive.rcfile-optimized-writer.enabled", "false") .put("hive.rcfile.writer.validate", "true") @@ -334,7 +332,6 @@ public void testExplicitPropertyMappings() .setDomainSocketPath("/foo") .setS3FileSystemType(S3FileSystemType.EMRFS) .setTextMaxLineLength(new DataSize(13, Unit.MEGABYTE)) - .setUseOrcColumnNames(true) .setAssumeCanonicalPartitionKeys(true) .setOrcDefaultBloomFilterFpp(0.96) .setRcfileOptimizedWriterEnabled(false) diff --git a/presto-hive/src/test/java/com/facebook/presto/hive/TestHiveDistributedNanQueries.java b/presto-hive/src/test/java/com/facebook/presto/hive/TestHiveDistributedNanQueries.java index fa2738c63fd72..fc74dbb2f4433 100644 --- a/presto-hive/src/test/java/com/facebook/presto/hive/TestHiveDistributedNanQueries.java +++ b/presto-hive/src/test/java/com/facebook/presto/hive/TestHiveDistributedNanQueries.java @@ -14,6 +14,7 @@ package com.facebook.presto.hive; +import com.facebook.presto.scalar.sql.SqlInvokedFunctionsPlugin; import com.facebook.presto.testing.QueryRunner; import com.facebook.presto.tests.AbstractTestNanQueries; import com.google.common.collect.ImmutableList; @@ -28,6 +29,9 @@ public class TestHiveDistributedNanQueries protected QueryRunner createQueryRunner() throws Exception { - return HiveQueryRunner.createQueryRunner(ImmutableList.of(), ImmutableMap.of("use-new-nan-definition", "true"), ImmutableMap.of(), Optional.empty()); + QueryRunner queryRunner = + HiveQueryRunner.createQueryRunner(ImmutableList.of(), ImmutableMap.of("use-new-nan-definition", "true"), ImmutableMap.of(), Optional.empty()); + queryRunner.installPlugin(new SqlInvokedFunctionsPlugin()); + return queryRunner; } } diff --git a/presto-hive/src/test/java/com/facebook/presto/hive/TestHiveDistributedQueries.java b/presto-hive/src/test/java/com/facebook/presto/hive/TestHiveDistributedQueries.java index c47da01ea5281..134cf4b101f1f 100644 --- a/presto-hive/src/test/java/com/facebook/presto/hive/TestHiveDistributedQueries.java +++ b/presto-hive/src/test/java/com/facebook/presto/hive/TestHiveDistributedQueries.java @@ -15,16 +15,22 @@ import com.facebook.presto.Session; import com.facebook.presto.hive.TestHiveEventListenerPlugin.TestingHiveEventListener; +import com.facebook.presto.scalar.sql.SqlInvokedFunctionsPlugin; import com.facebook.presto.spi.QueryId; import com.facebook.presto.spi.eventlistener.EventListener; import com.facebook.presto.testing.MaterializedResult; import com.facebook.presto.testing.QueryRunner; import com.facebook.presto.tests.AbstractTestDistributedQueries; import com.google.common.collect.ImmutableSet; +import com.google.common.io.Files; import org.intellij.lang.annotations.Language; import org.testng.annotations.AfterClass; import org.testng.annotations.Test; +import java.io.File; +import java.io.IOException; +import java.net.URI; +import java.net.URISyntaxException; import java.util.Set; import static com.facebook.presto.SystemSessionProperties.CTE_MATERIALIZATION_STRATEGY; @@ -33,8 +39,12 @@ import static com.facebook.presto.SystemSessionProperties.PUSHDOWN_SUBFIELDS_FOR_MAP_FUNCTIONS; import static com.facebook.presto.SystemSessionProperties.PUSHDOWN_SUBFIELDS_FROM_LAMBDA_ENABLED; import static com.facebook.presto.SystemSessionProperties.VERBOSE_OPTIMIZER_INFO_ENABLED; +import static com.facebook.presto.hive.HiveCommonSessionProperties.ORC_USE_COLUMN_NAMES; +import static com.facebook.presto.hive.HiveTestUtils.getHiveTableProperty; import static com.facebook.presto.sql.tree.ExplainType.Type.LOGICAL; import static com.google.common.collect.Iterables.getOnlyElement; +import static com.google.common.io.MoreFiles.deleteRecursively; +import static com.google.common.io.RecursiveDeleteOption.ALLOW_INSECURE; import static io.airlift.tpch.TpchTable.getTables; import static java.lang.String.format; import static java.util.stream.Collectors.joining; @@ -50,7 +60,9 @@ public class TestHiveDistributedQueries protected QueryRunner createQueryRunner() throws Exception { - return HiveQueryRunner.createQueryRunner(getTables()); + QueryRunner queryRunner = HiveQueryRunner.createQueryRunner(getTables()); + queryRunner.installPlugin(new SqlInvokedFunctionsPlugin()); + return queryRunner; } @Override @@ -290,5 +302,101 @@ public void testPushdownSubfieldForMapFunctions() } } + @Test + public void testOrcUseColumnNames() throws IOException, URISyntaxException + { + File externalTableDataDirectory = Files.createTempDir(); + String externalTableDataLocationUri = externalTableDataDirectory.toURI().toString(); + + try { + @Language("SQL") String createManagedTableSql = "" + + "create table test_orc_use_column_names (\n" + + " \"c1\" int,\n" + + " \"c2\" varchar\n" + + ")\n" + + "WITH (\n" + + " format = 'ORC'\n" + + ")"; + + assertUpdate(createManagedTableSql); + assertUpdate(format("insert into test_orc_use_column_names values (1, 'one')"), 1); + String tablePath = (String) getHiveTableProperty(getQueryRunner(), getSession(), "test_orc_use_column_names", (HiveTableLayoutHandle table) -> table.getTablePath()); + File managedTableDataDirectory = new File(new URI(tablePath).getRawPath()); + + assertTrue(managedTableDataDirectory.isDirectory(), "Source managed table data directory does not exist: " + managedTableDataDirectory); + File[] orcFiles = managedTableDataDirectory.listFiles(file -> file.isFile() + && !file.getName().contains(".")); + + if (orcFiles != null) { + for (File orcFile : orcFiles) { + File destinationFile = new File(externalTableDataDirectory, orcFile.getName()); + Files.copy(orcFile, destinationFile); + } + } + else { + throw new IllegalStateException("No ORC files found in managed table data directory: " + managedTableDataDirectory); + } + + @Language("SQL") String createMisMatchingExternalTableSql = format("" + + "CREATE TABLE test_orc_use_column_names_mismatching_ext (\n" + + " \"c2\" varchar,\n" + + " \"c1\" int\n" + + ")\n" + + "WITH (\n" + + " external_location = '%s',\n" + + " format = 'ORC'\n" + + ")", + externalTableDataLocationUri); + assertUpdate(createMisMatchingExternalTableSql); + + assertQueryFails(getSession(), + "select * from test_orc_use_column_names_mismatching_ext", + ".*java.io.IOException: Malformed ORC file. Can not read SQL type varchar from ORC stream .c1 of type INT.*"); + + Session useOrcColumnNamesSession = Session.builder(getQueryRunner().getDefaultSession()) + .setCatalogSessionProperty("hive", ORC_USE_COLUMN_NAMES, "true").build(); + assertQuerySucceeds(useOrcColumnNamesSession, "select * from test_orc_use_column_names_mismatching_ext"); + + @Language("SQL") String createDifferentColumnNameExternalTableSql = format("" + + "CREATE TABLE test_orc_use_column_names_different_column_name_ext (\n" + + " \"c1\" int,\n" + + " \"c3\" varchar\n" + + ")\n" + + "WITH (\n" + + " external_location = '%s',\n" + + " format = 'ORC'\n" + + ")", + externalTableDataLocationUri); + assertUpdate(createDifferentColumnNameExternalTableSql); + + assertQuery(useOrcColumnNamesSession, "select * from test_orc_use_column_names_different_column_name_ext", "VALUES (1, NULL)"); + + @Language("SQL") String createAdditionalColumnExternalTableSql = format("" + + "CREATE TABLE test_orc_use_column_names_additional_column_ext (\n" + + " \"c1\" int,\n" + + " \"c2\" varchar,\n" + + " \"c3\" varchar\n" + + ")\n" + + "WITH (\n" + + " external_location = '%s',\n" + + " format = 'ORC'\n" + + ")", + externalTableDataLocationUri); + assertUpdate(createAdditionalColumnExternalTableSql); + + assertQuery(useOrcColumnNamesSession, "select * from test_orc_use_column_names_additional_column_ext", "VALUES (1, 'one', NULL)"); + } + finally { + assertUpdate("DROP TABLE IF EXISTS test_orc_use_column_names"); + assertUpdate("DROP TABLE IF EXISTS test_orc_use_column_names_mismatching_ext"); + assertUpdate("DROP TABLE IF EXISTS test_orc_use_column_names_different_column_name_ext"); + assertUpdate("DROP TABLE IF EXISTS test_orc_use_column_names_additional_column_ext"); + + if (externalTableDataDirectory != null) { + deleteRecursively(externalTableDataDirectory.toPath(), ALLOW_INSECURE); + } + } + } + // Hive specific tests should normally go in TestHiveIntegrationSmokeTest } diff --git a/presto-hive/src/test/java/com/facebook/presto/hive/TestHiveDistributedQueriesWithExchangeMaterialization.java b/presto-hive/src/test/java/com/facebook/presto/hive/TestHiveDistributedQueriesWithExchangeMaterialization.java index 226bbf3888b53..a878243e32b91 100644 --- a/presto-hive/src/test/java/com/facebook/presto/hive/TestHiveDistributedQueriesWithExchangeMaterialization.java +++ b/presto-hive/src/test/java/com/facebook/presto/hive/TestHiveDistributedQueriesWithExchangeMaterialization.java @@ -14,6 +14,7 @@ package com.facebook.presto.hive; import com.facebook.presto.Session; +import com.facebook.presto.scalar.sql.SqlInvokedFunctionsPlugin; import com.facebook.presto.testing.MaterializedResult; import com.facebook.presto.testing.QueryRunner; import com.facebook.presto.tests.AbstractTestDistributedQueries; @@ -40,7 +41,9 @@ public class TestHiveDistributedQueriesWithExchangeMaterialization protected QueryRunner createQueryRunner() throws Exception { - return createMaterializingQueryRunner(getTables()); + QueryRunner queryRunner = createMaterializingQueryRunner(getTables()); + queryRunner.installPlugin(new SqlInvokedFunctionsPlugin()); + return queryRunner; } @Test diff --git a/presto-hive/src/test/java/com/facebook/presto/hive/TestHiveDistributedQueriesWithOptimizedRepartitioning.java b/presto-hive/src/test/java/com/facebook/presto/hive/TestHiveDistributedQueriesWithOptimizedRepartitioning.java index d0bf8f6499428..2477f9c31cf46 100644 --- a/presto-hive/src/test/java/com/facebook/presto/hive/TestHiveDistributedQueriesWithOptimizedRepartitioning.java +++ b/presto-hive/src/test/java/com/facebook/presto/hive/TestHiveDistributedQueriesWithOptimizedRepartitioning.java @@ -13,6 +13,7 @@ */ package com.facebook.presto.hive; +import com.facebook.presto.scalar.sql.SqlInvokedFunctionsPlugin; import com.facebook.presto.testing.QueryRunner; import com.facebook.presto.tests.AbstractTestDistributedQueries; import com.google.common.collect.ImmutableMap; @@ -30,13 +31,15 @@ public class TestHiveDistributedQueriesWithOptimizedRepartitioning protected QueryRunner createQueryRunner() throws Exception { - return HiveQueryRunner.createQueryRunner( + QueryRunner queryRunner = HiveQueryRunner.createQueryRunner( getTables(), ImmutableMap.of( "experimental.optimized-repartitioning", "true", // Use small SerializedPages to force flushing "driver.max-page-partitioning-buffer-size", "10000B"), Optional.empty()); + queryRunner.installPlugin(new SqlInvokedFunctionsPlugin()); + return queryRunner; } @Override diff --git a/presto-hive/src/test/java/com/facebook/presto/hive/TestHiveDistributedQueriesWithThriftRpc.java b/presto-hive/src/test/java/com/facebook/presto/hive/TestHiveDistributedQueriesWithThriftRpc.java index 6f877f35a4ce3..1bbc9590dcd50 100644 --- a/presto-hive/src/test/java/com/facebook/presto/hive/TestHiveDistributedQueriesWithThriftRpc.java +++ b/presto-hive/src/test/java/com/facebook/presto/hive/TestHiveDistributedQueriesWithThriftRpc.java @@ -13,6 +13,7 @@ */ package com.facebook.presto.hive; +import com.facebook.presto.scalar.sql.SqlInvokedFunctionsPlugin; import com.facebook.presto.testing.QueryRunner; import com.facebook.presto.tests.AbstractTestDistributedQueries; import com.google.common.collect.ImmutableMap; @@ -30,13 +31,15 @@ public class TestHiveDistributedQueriesWithThriftRpc protected QueryRunner createQueryRunner() throws Exception { - return HiveQueryRunner.createQueryRunner( + QueryRunner queryRunner = HiveQueryRunner.createQueryRunner( getTables(), ImmutableMap.of( "internal-communication.task-communication-protocol", "THRIFT", "internal-communication.server-info-communication-protocol", "THRIFT"), ImmutableMap.of(), Optional.empty()); + queryRunner.installPlugin(new SqlInvokedFunctionsPlugin()); + return queryRunner; } @Override diff --git a/presto-hive/src/test/java/com/facebook/presto/hive/TestHiveFileFormats.java b/presto-hive/src/test/java/com/facebook/presto/hive/TestHiveFileFormats.java index 8c117033a686b..f9fed4fdf2024 100644 --- a/presto-hive/src/test/java/com/facebook/presto/hive/TestHiveFileFormats.java +++ b/presto-hive/src/test/java/com/facebook/presto/hive/TestHiveFileFormats.java @@ -339,7 +339,7 @@ public void testOrc(int rowCount) assertThatFileFormat(ORC) .withColumns(TEST_COLUMNS) .withRowsCount(rowCount) - .isReadableByPageSource(new OrcBatchPageSourceFactory(FUNCTION_AND_TYPE_MANAGER, false, HDFS_ENVIRONMENT, STATS, 100, new StorageOrcFileTailSource(), StripeMetadataSourceFactory.of(new StorageStripeMetadataSource()))); + .isReadableByPageSource(new OrcBatchPageSourceFactory(FUNCTION_AND_TYPE_MANAGER, HDFS_ENVIRONMENT, STATS, 100, new StorageOrcFileTailSource(), StripeMetadataSourceFactory.of(new StorageStripeMetadataSource()))); } @Test(dataProvider = "rowCount") @@ -362,7 +362,7 @@ public void testOrcOptimizedWriter(int rowCount) .withSession(session) .withFileWriterFactory(new OrcFileWriterFactory(HDFS_ENVIRONMENT, new OutputStreamDataSinkFactory(), FUNCTION_AND_TYPE_MANAGER, new NodeVersion("test"), HIVE_STORAGE_TIME_ZONE, STATS, new OrcFileWriterConfig(), NO_ENCRYPTION)) .isReadableByRecordCursor(new GenericHiveRecordCursorProvider(HDFS_ENVIRONMENT)) - .isReadableByPageSource(new OrcBatchPageSourceFactory(FUNCTION_AND_TYPE_MANAGER, false, HDFS_ENVIRONMENT, STATS, 100, new StorageOrcFileTailSource(), StripeMetadataSourceFactory.of(new StorageStripeMetadataSource()))); + .isReadableByPageSource(new OrcBatchPageSourceFactory(FUNCTION_AND_TYPE_MANAGER, HDFS_ENVIRONMENT, STATS, 100, new StorageOrcFileTailSource(), StripeMetadataSourceFactory.of(new StorageStripeMetadataSource()))); } @Test(dataProvider = "rowCount") @@ -394,14 +394,15 @@ public void testOrcUseColumnNames(int rowCount) { TestingConnectorSession session = new TestingConnectorSession(getAllSessionProperties( new HiveClientConfig(), - new HiveCommonClientConfig())); + new HiveCommonClientConfig() + .setUseOrcColumnNames(true))); assertThatFileFormat(ORC) .withWriteColumns(TEST_COLUMNS) .withRowsCount(rowCount) .withReadColumns(Lists.reverse(TEST_COLUMNS)) .withSession(session) - .isReadableByPageSource(new OrcBatchPageSourceFactory(FUNCTION_AND_TYPE_MANAGER, true, HDFS_ENVIRONMENT, STATS, 100, new StorageOrcFileTailSource(), StripeMetadataSourceFactory.of(new StorageStripeMetadataSource()))); + .isReadableByPageSource(new OrcBatchPageSourceFactory(FUNCTION_AND_TYPE_MANAGER, HDFS_ENVIRONMENT, STATS, 100, new StorageOrcFileTailSource(), StripeMetadataSourceFactory.of(new StorageStripeMetadataSource()))); } @Test(dataProvider = "rowCount") @@ -419,7 +420,7 @@ public void testOrcUseColumnNamesCompatibility(int rowCount) .withRowsCount(rowCount) .withReadColumns(TEST_COLUMNS) .withSession(session) - .isReadableByPageSource(new OrcBatchPageSourceFactory(FUNCTION_AND_TYPE_MANAGER, true, HDFS_ENVIRONMENT, STATS, 100, new StorageOrcFileTailSource(), StripeMetadataSourceFactory.of(new StorageStripeMetadataSource()))); + .isReadableByPageSource(new OrcBatchPageSourceFactory(FUNCTION_AND_TYPE_MANAGER, HDFS_ENVIRONMENT, STATS, 100, new StorageOrcFileTailSource(), StripeMetadataSourceFactory.of(new StorageStripeMetadataSource()))); } private static List getHiveColumnNameColumns() @@ -628,7 +629,7 @@ public void testTruncateVarcharColumn() assertThatFileFormat(ORC) .withWriteColumns(ImmutableList.of(writeColumn)) .withReadColumns(ImmutableList.of(readColumn)) - .isReadableByPageSource(new OrcBatchPageSourceFactory(FUNCTION_AND_TYPE_MANAGER, false, HDFS_ENVIRONMENT, STATS, 100, new StorageOrcFileTailSource(), StripeMetadataSourceFactory.of(new StorageStripeMetadataSource()))); + .isReadableByPageSource(new OrcBatchPageSourceFactory(FUNCTION_AND_TYPE_MANAGER, HDFS_ENVIRONMENT, STATS, 100, new StorageOrcFileTailSource(), StripeMetadataSourceFactory.of(new StorageStripeMetadataSource()))); assertThatFileFormat(PARQUET) .withWriteColumns(ImmutableList.of(writeColumn)) @@ -676,7 +677,7 @@ public void testFailForLongVarcharPartitionColumn() assertThatFileFormat(ORC) .withColumns(columns) - .isFailingForPageSource(new OrcBatchPageSourceFactory(FUNCTION_AND_TYPE_MANAGER, false, HDFS_ENVIRONMENT, STATS, 100, new StorageOrcFileTailSource(), StripeMetadataSourceFactory.of(new StorageStripeMetadataSource())), expectedErrorCode, expectedMessage); + .isFailingForPageSource(new OrcBatchPageSourceFactory(FUNCTION_AND_TYPE_MANAGER, HDFS_ENVIRONMENT, STATS, 100, new StorageOrcFileTailSource(), StripeMetadataSourceFactory.of(new StorageStripeMetadataSource())), expectedErrorCode, expectedMessage); assertThatFileFormat(PARQUET) .withColumns(columns) diff --git a/presto-hive/src/test/java/com/facebook/presto/hive/TestHiveIntegrationSmokeTest.java b/presto-hive/src/test/java/com/facebook/presto/hive/TestHiveIntegrationSmokeTest.java index 1f765eab22008..51a0c7807eb22 100644 --- a/presto-hive/src/test/java/com/facebook/presto/hive/TestHiveIntegrationSmokeTest.java +++ b/presto-hive/src/test/java/com/facebook/presto/hive/TestHiveIntegrationSmokeTest.java @@ -23,12 +23,10 @@ import com.facebook.presto.hive.HiveClientConfig.InsertExistingPartitionsBehavior; import com.facebook.presto.metadata.InsertTableHandle; import com.facebook.presto.metadata.Metadata; -import com.facebook.presto.metadata.TableLayout; import com.facebook.presto.spi.CatalogSchemaTableName; import com.facebook.presto.spi.ColumnMetadata; import com.facebook.presto.spi.ConnectorId; import com.facebook.presto.spi.ConnectorSession; -import com.facebook.presto.spi.Constraint; import com.facebook.presto.spi.TableHandle; import com.facebook.presto.spi.TableMetadata; import com.facebook.presto.spi.plan.MarkDistinctNode; @@ -75,7 +73,6 @@ import java.util.Optional; import java.util.function.BiConsumer; import java.util.function.Consumer; -import java.util.function.Function; import java.util.stream.LongStream; import static com.facebook.airlift.json.JsonCodec.jsonCodec; @@ -130,6 +127,7 @@ import static com.facebook.presto.hive.HiveTableProperties.PARTITIONED_BY_PROPERTY; import static com.facebook.presto.hive.HiveTableProperties.STORAGE_FORMAT_PROPERTY; import static com.facebook.presto.hive.HiveTestUtils.FUNCTION_AND_TYPE_MANAGER; +import static com.facebook.presto.hive.HiveTestUtils.getHiveTableProperty; import static com.facebook.presto.hive.HiveUtil.columnExtraInfo; import static com.facebook.presto.spi.security.SelectedRole.Type.ROLE; import static com.facebook.presto.sql.analyzer.FeaturesConfig.JoinDistributionType.BROADCAST; @@ -2083,31 +2081,14 @@ private TableMetadata getTableMetadata(String catalog, String schema, String tab }); } - private Object getHiveTableProperty(String tableName, Function propertyGetter) - { - Session session = getSession(); - Metadata metadata = ((DistributedQueryRunner) getQueryRunner()).getCoordinator().getMetadata(); - - return transaction(getQueryRunner().getTransactionManager(), getQueryRunner().getAccessControl()) - .readOnly() - .execute(session, transactionSession -> { - Optional tableHandle = metadata.getMetadataResolver(transactionSession).getTableHandle(new QualifiedObjectName(catalog, TPCH_SCHEMA, tableName)); - assertTrue(tableHandle.isPresent()); - - TableLayout layout = metadata.getLayout(transactionSession, tableHandle.get(), Constraint.alwaysTrue(), Optional.empty()) - .getLayout(); - return propertyGetter.apply((HiveTableLayoutHandle) layout.getNewTableHandle().getLayout().get()); - }); - } - private List getPartitions(String tableName) { - return (List) getHiveTableProperty(tableName, (HiveTableLayoutHandle table) -> getPartitions(table)); + return (List) getHiveTableProperty(getQueryRunner(), getSession(), tableName, (HiveTableLayoutHandle table) -> getPartitions(table)); } private int getBucketCount(String tableName) { - return (int) getHiveTableProperty(tableName, (HiveTableLayoutHandle table) -> table.getBucketHandle().get().getTableBucketCount()); + return (int) getHiveTableProperty(getQueryRunner(), getSession(), tableName, (HiveTableLayoutHandle table) -> table.getBucketHandle().get().getTableBucketCount()); } @Test diff --git a/presto-hive/src/test/java/com/facebook/presto/hive/TestHiveLogicalPlanner.java b/presto-hive/src/test/java/com/facebook/presto/hive/TestHiveLogicalPlanner.java index 07d0e3a7c6385..534beafe5a0b6 100644 --- a/presto-hive/src/test/java/com/facebook/presto/hive/TestHiveLogicalPlanner.java +++ b/presto-hive/src/test/java/com/facebook/presto/hive/TestHiveLogicalPlanner.java @@ -85,6 +85,7 @@ import static com.facebook.presto.SystemSessionProperties.PUSHDOWN_DEREFERENCE_ENABLED; import static com.facebook.presto.SystemSessionProperties.PUSHDOWN_SUBFIELDS_ENABLED; import static com.facebook.presto.SystemSessionProperties.PUSHDOWN_SUBFIELDS_FOR_MAP_FUNCTIONS; +import static com.facebook.presto.SystemSessionProperties.UTILIZE_UNIQUE_PROPERTY_IN_QUERY_PLANNING; import static com.facebook.presto.common.function.OperatorType.EQUAL; import static com.facebook.presto.common.predicate.Domain.create; import static com.facebook.presto.common.predicate.Domain.multipleValues; @@ -130,6 +131,7 @@ import static com.facebook.presto.sql.planner.assertions.PlanMatchPattern.node; import static com.facebook.presto.sql.planner.assertions.PlanMatchPattern.output; import static com.facebook.presto.sql.planner.assertions.PlanMatchPattern.project; +import static com.facebook.presto.sql.planner.assertions.PlanMatchPattern.semiJoin; import static com.facebook.presto.sql.planner.assertions.PlanMatchPattern.strictTableScan; import static com.facebook.presto.sql.planner.assertions.PlanMatchPattern.values; import static com.facebook.presto.sql.planner.optimizations.PlanNodeSearcher.searchFrom; @@ -2153,6 +2155,336 @@ public void testPartialAggregatePushdown() } } + @Test + public void testRowId() + { + Session session = Session.builder(getQueryRunner().getDefaultSession()) + .setSystemProperty(UTILIZE_UNIQUE_PROPERTY_IN_QUERY_PLANNING, "true") + .build(); + String sql; + sql = "SELECT\n" + + " unique_id AS unique_id,\n" + + " ARBITRARY(name) AS customer_name,\n" + + " ARRAY_AGG(orderkey) AS orders_info\n" + + "FROM (\n" + + " SELECT\n" + + " customer.name,\n" + + " orders.orderkey,\n" + + " customer.\"$row_id\" AS unique_id\n" + + " FROM customer\n" + + " LEFT JOIN orders\n" + + " ON customer.custkey = orders.custkey\n" + + " AND orders.orderstatus IN ('O', 'F')\n" + + " AND orders.orderdate BETWEEN DATE '1995-01-01' AND DATE '1995-12-31'\n" + + " WHERE\n" + + " customer.nationkey IN (1, 2, 3, 4, 5)\n" + + ")\n" + + "GROUP BY\n" + + " unique_id"; + + assertPlan(session, + sql, + anyTree( + aggregation( + ImmutableMap.of(), + join( + anyTree(tableScan("customer")), + anyTree(tableScan("orders")))))); + + sql = "select \"$row_id\", count(*) from orders group by 1"; + assertPlan(sql, anyTree(aggregation(ImmutableMap.of(), tableScan("orders")))); + sql = "SELECT\n" + + " customer.\"$row_id\" AS unique_id,\n" + + " COUNT(orders.orderkey) AS order_count,\n" + + " ARRAY_AGG(orders.orderkey) AS order_keys\n" + + "FROM customer\n" + + "LEFT JOIN orders\n" + + " ON customer.custkey = orders.custkey\n" + + "WHERE customer.acctbal > 1000\n" + + "GROUP BY customer.\"$row_id\""; + assertPlan(sql, anyTree( + aggregation(ImmutableMap.of(), + join( + anyTree(tableScan("customer")), + anyTree(tableScan("orders")))))); + + sql = "SELECT\n" + + " unique_id,\n" + + " MAX(orderdate) AS latest_order_date\n" + + "FROM (\n" + + " SELECT\n" + + " customer.\"$row_id\" AS unique_id,\n" + + " orders.orderdate\n" + + " FROM customer\n" + + " JOIN orders\n" + + " ON customer.custkey = orders.custkey\n" + + " WHERE orders.orderstatus = 'O'\n" + + ") t\n" + + "GROUP BY unique_id"; + assertPlan(sql, anyTree( + aggregation(ImmutableMap.of(), + join( + anyTree(tableScan("customer")), + anyTree(tableScan("orders")))))); + + sql = "SELECT\n" + + " DISTINCT customer.\"$row_id\" AS unique_id,\n" + + " customer.name\n" + + "FROM customer\n" + + "WHERE customer.nationkey = 1"; + + assertPlan(sql, anyTree( + aggregation(ImmutableMap.of(), + project(filter(tableScan("customer")))))); + + sql = "SELECT\n" + + " c.name,\n" + + " o.orderkey,\n" + + " c.\"$row_id\" AS customer_row_id,\n" + + " o.\"$row_id\" AS order_row_id\n" + + "FROM customer c\n" + + "JOIN orders o\n" + + " ON c.\"$row_id\" = o.\"$row_id\"\n" + + "WHERE o.totalprice > 10000"; + assertPlan(sql, anyTree( + join( + exchange(anyTree(tableScan("customer"))), + exchange(anyTree(tableScan("orders")))))); + + sql = "SELECT\n" + + " \"$row_id\" AS unique_id,\n" + + " orderstatus,\n" + + " COUNT(*) AS cnt\n" + + "FROM orders\n" + + "GROUP BY \"$row_id\", orderstatus"; + assertPlan(sql, anyTree( + aggregation(ImmutableMap.of(), + project(tableScan("orders"))))); + + sql = "SELECT\n" + + " o1.orderkey,\n" + + " o2.totalprice\n" + + "FROM orders o1\n" + + "JOIN orders o2\n" + + " ON o1.\"$row_id\" = o2.\"$row_id\"\n" + + "WHERE o1.orderstatus = 'O'"; + assertPlan(sql, anyTree( + join( + exchange(anyTree(tableScan("orders"))), + exchange(anyTree(tableScan("orders")))))); + + sql = "SELECT\n" + + " orderkey,\n" + + " totalprice\n" + + "FROM orders o1\n" + + "WHERE EXISTS (\n" + + " SELECT 1\n" + + " FROM orders o2\n" + + " WHERE o1.\"$row_id\" = o2.\"$row_id\"\n" + + " AND o2.orderstatus = 'F'\n" + + ")"; + assertPlan(sql, anyTree( + join( + exchange(anyTree(tableScan("orders"))), + exchange(anyTree(aggregation(ImmutableMap.of(), project(filter(tableScan("orders"))))))))); + + sql = "SELECT\n" + + " custkey,\n" + + " name\n" + + "FROM customer\n" + + "WHERE \"$row_id\" IN (\n" + + " SELECT \"$row_id\"\n" + + " FROM customer\n" + + " WHERE acctbal > 5000\n" + + ")"; + assertPlan(sql, anyTree( + semiJoin( + anyTree(tableScan("customer")), + anyTree(tableScan("customer"))))); + + sql = "SELECT \"$row_id\", orderkey FROM orders WHERE orderstatus = 'O'\n" + + "UNION\n" + + "SELECT \"$row_id\", orderkey FROM orders WHERE orderstatus = 'F'"; + assertPlan(sql, anyTree( + aggregation(ImmutableMap.of(), + project(filter(tableScan("orders")))))); + + sql = "SELECT\n" + + " c.\"$row_id\",\n" + + " COUNT(o.orderkey) AS order_count,\n" + + " SUM(o.totalprice) AS total_spent,\n" + + " MAX(o.orderdate) AS latest_order,\n" + + " MIN(o.orderdate) AS first_order\n" + + "FROM customer c\n" + + "LEFT JOIN orders o ON c.custkey = o.custkey\n" + + "GROUP BY c.\"$row_id\""; + assertPlan(sql, anyTree( + aggregation(ImmutableMap.of(), + join( + anyTree(tableScan("customer")), + anyTree(tableScan("orders")))))); + + sql = "SELECT\n" + + " \"$row_id\",\n" + + " COUNT(*) AS cnt\n" + + "FROM orders\n" + + "GROUP BY \"$row_id\"\n" + + "HAVING COUNT(*) = 1"; + assertPlan(sql, anyTree( + project(filter( + aggregation(ImmutableMap.of(), + tableScan("orders")))))); + + sql = "SELECT\n" + + " c.custkey,\n" + + " c.name,\n" + + " (\n" + + " SELECT COUNT(*)\n" + + " FROM orders o\n" + + " WHERE o.custkey = c.custkey\n" + + " AND o.\"$row_id\" IS NOT NULL\n" + + " ) AS order_count\n" + + "FROM customer c"; + assertPlan(sql, anyTree( + project( + join( + anyTree(tableScan("customer")), + anyTree(tableScan("orders")))))); + + sql = "SELECT\n" + + " outer_query.customer_id,\n" + + " outer_query.order_count\n" + + "FROM (\n" + + " SELECT\n" + + " c.\"$row_id\" AS customer_id,\n" + + " COUNT(DISTINCT o.\"$row_id\") AS order_count\n" + + " FROM customer c\n" + + " LEFT JOIN orders o ON c.custkey = o.custkey\n" + + " WHERE c.nationkey IN (1, 2, 3)\n" + + " GROUP BY c.\"$row_id\"\n" + + ") outer_query\n" + + "WHERE outer_query.order_count > 2"; + assertPlan(sql, anyTree( + aggregation(ImmutableMap.of(), + project( + join( + anyTree(tableScan("customer")), + anyTree(tableScan("orders"))))))); + + sql = "SELECT\n" + + " c.custkey,\n" + + " c.name\n" + + "FROM customer c\n" + + "WHERE NOT EXISTS (\n" + + " SELECT 1\n" + + " FROM orders o\n" + + " WHERE c.\"$row_id\" = o.\"$row_id\"\n" + + ")"; + assertPlan(sql, anyTree( + join( + exchange(anyTree(tableScan("customer"))), + exchange(anyTree(aggregation(ImmutableMap.of(), tableScan("orders"))))))); + + sql = "SELECT\n" + + " c.name,\n" + + " o.orderkey,\n" + + " l.linenumber\n" + + "FROM customer c\n" + + "JOIN orders o ON c.custkey = o.custkey\n" + + "JOIN lineitem l ON o.orderkey = l.orderkey\n" + + "WHERE c.\"$row_id\" IS NOT NULL\n" + + " AND o.\"$row_id\" IS NOT NULL\n" + + " AND l.\"$row_id\" IS NOT NULL"; + assertPlan(sql, anyTree( + join( + anyTree( + join( + anyTree(tableScan("customer")), + anyTree(tableScan("orders")))), + anyTree(tableScan("lineitem"))))); + + sql = "SELECT\n" + + " c.\"$row_id\" AS customer_row_id,\n" + + " o.\"$row_id\" AS order_row_id,\n" + + " c.name,\n" + + " o.orderkey\n" + + "FROM customer c\n" + + "CROSS JOIN orders o\n" + + "WHERE c.nationkey = 1 AND o.orderstatus = 'O'"; + assertPlan(sql, anyTree( + join( + anyTree(tableScan("customer")), + anyTree(tableScan("orders"))))); + + sql = "SELECT\n" + + " orderstatus,\n" + + " COUNT(\"$row_id\") AS row_count,\n" + + " MIN(\"$row_id\") AS min_row_id\n" + + "FROM orders\n" + + "GROUP BY orderstatus"; + assertPlan(sql, anyTree( + aggregation(ImmutableMap.of(), + exchange(anyTree(tableScan("orders")))))); + + sql = "SELECT\n" + + " xxhash64(\"$row_id\") AS bucket,\n" + + " COUNT(*) AS cnt\n" + + "FROM orders\n" + + "GROUP BY xxhash64(\"$row_id\")"; + assertPlan(sql, anyTree( + aggregation(ImmutableMap.of(), + exchange(anyTree(tableScan("orders")))))); + + sql = "SELECT \"$row_id\", 'customer' AS source FROM customer\n" + + "UNION ALL\n" + + "SELECT \"$row_id\", 'orders' AS source FROM orders"; + assertPlan(sql, anyTree( + exchange( + anyTree(tableScan("customer")), + anyTree(tableScan("orders"))))); + + sql = "SELECT\n" + + " o.\"$row_id\" AS order_row_id,\n" + + " o.orderkey,\n" + + " l.linenumber\n" + + "FROM orders o\n" + + "JOIN lineitem l ON o.orderkey = l.orderkey\n" + + "WHERE o.orderstatus = 'O'"; + assertPlan(sql, anyTree( + join( + exchange(anyTree(tableScan("lineitem"))), + exchange(anyTree(tableScan("orders")))))); + + sql = "SELECT\n" + + " orderkey,\n" + + " totalprice\n" + + "FROM orders o1\n" + + "WHERE \"$row_id\" > (\n" + + " SELECT MIN(\"$row_id\")\n" + + " FROM orders o2\n" + + " WHERE o2.orderstatus = 'F'\n" + + ")"; + assertPlan(sql, anyTree( + join( + tableScan("orders"), + exchange(anyTree(tableScan("orders")))))); + + sql = "SELECT\n" + + " CASE \n" + + " WHEN nationkey = 1 THEN \"$row_id\"\n" + + " ELSE cast('default' as varbinary)\n" + + " END AS conditional_row_id,\n" + + " COUNT(*) AS cnt\n" + + "FROM customer\n" + + "GROUP BY CASE \n" + + " WHEN nationkey = 1 THEN \"$row_id\"\n" + + " ELSE cast('default' as varbinary)\n" + + " END"; + assertPlan(sql, anyTree( + aggregation(ImmutableMap.of(), + exchange(anyTree(tableScan("customer")))))); + } + private static Set toSubfields(String... subfieldPaths) { return Arrays.stream(subfieldPaths) @@ -2170,6 +2502,11 @@ private void assertPushdownSubfields(Session session, String query, String table assertPlan(session, query, anyTree(tableScan(tableName, requiredSubfields))); } + private static PlanMatchPattern tableScan(String expectedTableName) + { + return PlanMatchPattern.tableScan(expectedTableName); + } + private static PlanMatchPattern tableScan(String expectedTableName, Map> expectedRequiredSubfields) { return PlanMatchPattern.tableScan(expectedTableName).with(new HiveTableScanMatcher(expectedRequiredSubfields)); diff --git a/presto-hive/src/test/java/com/facebook/presto/hive/TestHivePushdownDistributedQueries.java b/presto-hive/src/test/java/com/facebook/presto/hive/TestHivePushdownDistributedQueries.java index 28ec3c9ef3491..dc816175ffb12 100644 --- a/presto-hive/src/test/java/com/facebook/presto/hive/TestHivePushdownDistributedQueries.java +++ b/presto-hive/src/test/java/com/facebook/presto/hive/TestHivePushdownDistributedQueries.java @@ -13,6 +13,7 @@ */ package com.facebook.presto.hive; +import com.facebook.presto.scalar.sql.SqlInvokedFunctionsPlugin; import com.facebook.presto.testing.MaterializedResult; import com.facebook.presto.testing.QueryRunner; import com.facebook.presto.tests.AbstractTestDistributedQueries; @@ -34,7 +35,7 @@ public class TestHivePushdownDistributedQueries protected QueryRunner createQueryRunner() throws Exception { - return HiveQueryRunner.createQueryRunner( + QueryRunner queryRunner = HiveQueryRunner.createQueryRunner( getTables(), ImmutableMap.of("experimental.pushdown-subfields-enabled", "true", "experimental.pushdown-dereference-enabled", "true"), @@ -44,6 +45,8 @@ protected QueryRunner createQueryRunner() "hive.partial_aggregation_pushdown_enabled", "true", "hive.partial_aggregation_pushdown_for_variable_length_datatypes_enabled", "true"), Optional.empty()); + queryRunner.installPlugin(new SqlInvokedFunctionsPlugin()); + return queryRunner; } @Override diff --git a/presto-hive/src/test/java/com/facebook/presto/hive/TestLambdaSubfieldPruning.java b/presto-hive/src/test/java/com/facebook/presto/hive/TestLambdaSubfieldPruning.java index a0559b03b0302..b99eff80743ab 100644 --- a/presto-hive/src/test/java/com/facebook/presto/hive/TestLambdaSubfieldPruning.java +++ b/presto-hive/src/test/java/com/facebook/presto/hive/TestLambdaSubfieldPruning.java @@ -14,6 +14,7 @@ package com.facebook.presto.hive; import com.facebook.presto.Session; +import com.facebook.presto.scalar.sql.SqlInvokedFunctionsPlugin; import com.facebook.presto.testing.QueryRunner; import com.facebook.presto.tests.AbstractTestQueryFramework; import com.facebook.presto.tests.DistributedQueryRunner; @@ -126,6 +127,7 @@ private static DistributedQueryRunner createLineItemExTable(DistributedQueryRunn "FROM lineitem \n"); } + queryRunner.installPlugin(new SqlInvokedFunctionsPlugin()); return queryRunner; } diff --git a/presto-hive/src/test/java/com/facebook/presto/hive/TestOrcBatchPageSourceMemoryTracking.java b/presto-hive/src/test/java/com/facebook/presto/hive/TestOrcBatchPageSourceMemoryTracking.java index 42509b0089f7b..6144871f10040 100644 --- a/presto-hive/src/test/java/com/facebook/presto/hive/TestOrcBatchPageSourceMemoryTracking.java +++ b/presto-hive/src/test/java/com/facebook/presto/hive/TestOrcBatchPageSourceMemoryTracking.java @@ -456,7 +456,6 @@ public ConnectorPageSource newPageSource(FileFormatDataSourceStats stats, Connec OrcBatchPageSourceFactory orcPageSourceFactory = new OrcBatchPageSourceFactory( FUNCTION_AND_TYPE_MANAGER, - false, HDFS_ENVIRONMENT, stats, 100, diff --git a/presto-hive/src/test/java/com/facebook/presto/hive/TestParquetDistributedQueries.java b/presto-hive/src/test/java/com/facebook/presto/hive/TestParquetDistributedQueries.java index d8bb271cebcd0..0aef10a5b4bf0 100644 --- a/presto-hive/src/test/java/com/facebook/presto/hive/TestParquetDistributedQueries.java +++ b/presto-hive/src/test/java/com/facebook/presto/hive/TestParquetDistributedQueries.java @@ -14,6 +14,7 @@ package com.facebook.presto.hive; import com.facebook.presto.Session; +import com.facebook.presto.scalar.sql.SqlInvokedFunctionsPlugin; import com.facebook.presto.testing.MaterializedResult; import com.facebook.presto.testing.QueryRunner; import com.facebook.presto.tests.AbstractTestDistributedQueries; @@ -46,7 +47,7 @@ protected QueryRunner createQueryRunner() .put("hive.partial_aggregation_pushdown_enabled", "true") .put("hive.partial_aggregation_pushdown_for_variable_length_datatypes_enabled", "true") .build(); - return HiveQueryRunner.createQueryRunner( + QueryRunner queryRunner = HiveQueryRunner.createQueryRunner( getTables(), ImmutableMap.of( "experimental.pushdown-subfields-enabled", "true", @@ -54,6 +55,8 @@ protected QueryRunner createQueryRunner() "sql-standard", parquetProperties, Optional.empty()); + queryRunner.installPlugin(new SqlInvokedFunctionsPlugin()); + return queryRunner; } @Test diff --git a/presto-hive/src/test/java/com/facebook/presto/hive/benchmark/FileFormat.java b/presto-hive/src/test/java/com/facebook/presto/hive/benchmark/FileFormat.java index 11b67fea8b68e..eea2d2a876fd3 100644 --- a/presto-hive/src/test/java/com/facebook/presto/hive/benchmark/FileFormat.java +++ b/presto-hive/src/test/java/com/facebook/presto/hive/benchmark/FileFormat.java @@ -161,7 +161,6 @@ public ConnectorPageSource createFileFormatReader(ConnectorSession session, Hdfs { HiveBatchPageSourceFactory pageSourceFactory = new OrcBatchPageSourceFactory( FUNCTION_AND_TYPE_MANAGER, - false, hdfsEnvironment, new FileFormatDataSourceStats(), 100, diff --git a/presto-hudi/pom.xml b/presto-hudi/pom.xml index 63c0ced3bd811..50546a3f996f6 100644 --- a/presto-hudi/pom.xml +++ b/presto-hudi/pom.xml @@ -148,6 +148,8 @@ org.apache.hudi hudi-presto-bundle + + 1.0.2 diff --git a/presto-hudi/src/main/java/com/facebook/presto/hudi/HudiMetadata.java b/presto-hudi/src/main/java/com/facebook/presto/hudi/HudiMetadata.java index 34d2676d86015..ac6018c3f142c 100644 --- a/presto-hudi/src/main/java/com/facebook/presto/hudi/HudiMetadata.java +++ b/presto-hudi/src/main/java/com/facebook/presto/hudi/HudiMetadata.java @@ -101,6 +101,7 @@ public ConnectorTableHandle getTableHandle(ConnectorSession session, SchemaTable } return new HudiTableHandle( + Optional.of(table), table.getDatabaseName(), table.getTableName(), table.getStorage().getLocation(), @@ -115,7 +116,11 @@ public Optional getSystemTable(ConnectorSession session, SchemaTabl } @Override - public List getTableLayouts(ConnectorSession session, ConnectorTableHandle tableHandle, Constraint constraint, Optional> desiredColumns) + public ConnectorTableLayoutResult getTableLayoutForConstraint( + ConnectorSession session, + ConnectorTableHandle tableHandle, + Constraint constraint, + Optional> desiredColumns) { HudiTableHandle handle = (HudiTableHandle) tableHandle; Table table = getTable(session, tableHandle); @@ -127,7 +132,7 @@ public List getTableLayouts(ConnectorSession session partitionColumns, table.getParameters(), constraint.getSummary())); - return ImmutableList.of(new ConnectorTableLayoutResult(layout, constraint.getSummary())); + return new ConnectorTableLayoutResult(layout, constraint.getSummary()); } @Override @@ -241,7 +246,7 @@ static List fromPartitionColumns(List partitionColumns return builder.build(); } - static List fromDataColumns(List dataColumns) + public static List fromDataColumns(List dataColumns) { ImmutableList.Builder builder = ImmutableList.builder(); int id = 0; diff --git a/presto-hudi/src/main/java/com/facebook/presto/hudi/HudiPartitionManager.java b/presto-hudi/src/main/java/com/facebook/presto/hudi/HudiPartitionManager.java index ba80fd2f516c7..0c1451d89479b 100644 --- a/presto-hudi/src/main/java/com/facebook/presto/hudi/HudiPartitionManager.java +++ b/presto-hudi/src/main/java/com/facebook/presto/hudi/HudiPartitionManager.java @@ -14,6 +14,7 @@ package com.facebook.presto.hudi; +import com.facebook.airlift.log.Logger; import com.facebook.presto.common.predicate.Domain; import com.facebook.presto.common.predicate.NullableValue; import com.facebook.presto.common.predicate.TupleDomain; @@ -23,6 +24,8 @@ import com.facebook.presto.hive.metastore.Column; import com.facebook.presto.hive.metastore.ExtendedHiveMetastore; import com.facebook.presto.hive.metastore.MetastoreContext; +import com.facebook.presto.hive.metastore.Partition; +import com.facebook.presto.hive.metastore.StorageFormat; import com.facebook.presto.hive.metastore.Table; import com.facebook.presto.spi.ColumnHandle; import com.facebook.presto.spi.ConnectorSession; @@ -45,10 +48,10 @@ import static com.facebook.presto.hudi.HudiMetadata.fromPartitionColumns; import static com.facebook.presto.hudi.HudiMetadata.toMetastoreContext; import static java.util.Objects.requireNonNull; -import static java.util.stream.Collectors.toList; public class HudiPartitionManager { + private static final Logger log = Logger.get(HudiPartitionManager.class); private final TypeManager typeManager; @Inject @@ -57,10 +60,11 @@ public HudiPartitionManager(TypeManager typeManager) this.typeManager = requireNonNull(typeManager, "typeManager is null"); } - public List getEffectivePartitions( + public Map getEffectivePartitions( ConnectorSession connectorSession, ExtendedHiveMetastore metastore, SchemaTableName schemaTableName, + String tableBasePath, TupleDomain constraintSummary) { MetastoreContext metastoreContext = toMetastoreContext(connectorSession); @@ -68,7 +72,17 @@ public List getEffectivePartitions( Verify.verify(table.isPresent()); List partitionColumns = table.get().getPartitionColumns(); if (partitionColumns.isEmpty()) { - return ImmutableList.of(""); + return ImmutableMap.of( + "", Partition.builder() + .setCatalogName(Optional.empty()) + .setDatabaseName(schemaTableName.getSchemaName()) + .setTableName(schemaTableName.getTableName()) + .withStorage(storageBuilder -> + storageBuilder.setLocation(tableBasePath) + .setStorageFormat(StorageFormat.VIEW_STORAGE_FORMAT)) + .setColumns(ImmutableList.of()) + .setValues(ImmutableList.of()) + .build()); } Map partitionPredicate = new HashMap<>(); @@ -87,17 +101,29 @@ public List getEffectivePartitions( List partitionNames = metastore.getPartitionNamesByFilter(metastoreContext, schemaTableName.getSchemaName(), schemaTableName.getTableName(), partitionPredicate); List partitionTypes = partitionColumns.stream() .map(column -> typeManager.getType(column.getType().getTypeSignature())) - .collect(toList()); + .toList(); - return partitionNames.stream() - .map(PartitionNameWithVersion::getPartitionName) + List filteredPartitionNames = partitionNames.stream() // Apply extra filters which could not be done by getPartitionNamesByFilter, similar to filtering in HivePartitionManager#getPartitionsIterator - .filter(partitionName -> parseValuesAndFilterPartition( - partitionName, + .filter(partitionNameWithVersion -> parseValuesAndFilterPartition( + partitionNameWithVersion.getPartitionName(), hudiColumnHandles, partitionTypes, constraintSummary)) - .collect(toList()); + .toList(); + Map> partitionsByNames = metastore.getPartitionsByNames(metastoreContext, schemaTableName.getSchemaName(), schemaTableName.getTableName(), filteredPartitionNames); + List partitionsNotFound = partitionsByNames.entrySet().stream().filter(e -> e.getValue().isEmpty()).map(Map.Entry::getKey).toList(); + if (!partitionsNotFound.isEmpty()) { + log.warn("The following partitions were not found in the metastore for table {}.{}: {}", + schemaTableName.getSchemaName(), + schemaTableName.getTableName(), + partitionsNotFound); + } + + return partitionsByNames + .entrySet().stream() + .filter(e -> e.getValue().isPresent()) + .collect(Collectors.toMap(Map.Entry::getKey, e -> e.getValue().get())); } private boolean parseValuesAndFilterPartition( diff --git a/presto-hudi/src/main/java/com/facebook/presto/hudi/HudiRecordCursors.java b/presto-hudi/src/main/java/com/facebook/presto/hudi/HudiRecordCursors.java index dd73e684fc370..fbab341ade696 100644 --- a/presto-hudi/src/main/java/com/facebook/presto/hudi/HudiRecordCursors.java +++ b/presto-hudi/src/main/java/com/facebook/presto/hudi/HudiRecordCursors.java @@ -55,6 +55,7 @@ import static org.apache.hadoop.hive.serde2.ColumnProjectionUtils.READ_ALL_COLUMNS; import static org.apache.hadoop.hive.serde2.ColumnProjectionUtils.READ_COLUMN_IDS_CONF_STR; import static org.apache.hadoop.hive.serde2.ColumnProjectionUtils.READ_COLUMN_NAMES_CONF_STR; +import static org.apache.hudi.common.config.HoodieReaderConfig.FILE_GROUP_READER_ENABLED; class HudiRecordCursors { @@ -105,6 +106,7 @@ public static RecordCursor createRealtimeRecordCursor( jobConf.setBoolean(READ_ALL_COLUMNS, false); jobConf.set(READ_COLUMN_IDS_CONF_STR, join(dataColumns, HudiColumnHandle::getId)); jobConf.set(READ_COLUMN_NAMES_CONF_STR, join(dataColumns, HudiColumnHandle::getName)); + jobConf.setBoolean(FILE_GROUP_READER_ENABLED.key(), false); schema.stringPropertyNames() .forEach(name -> jobConf.set(name, schema.getProperty(name))); refineCompressionCodecs(jobConf); diff --git a/presto-hudi/src/main/java/com/facebook/presto/hudi/HudiSplitManager.java b/presto-hudi/src/main/java/com/facebook/presto/hudi/HudiSplitManager.java index e1f38c291aa88..24d14fbaa74dc 100644 --- a/presto-hudi/src/main/java/com/facebook/presto/hudi/HudiSplitManager.java +++ b/presto-hudi/src/main/java/com/facebook/presto/hudi/HudiSplitManager.java @@ -19,9 +19,7 @@ import com.facebook.presto.hive.HdfsEnvironment; import com.facebook.presto.hive.filesystem.ExtendedFileSystem; import com.facebook.presto.hive.metastore.ExtendedHiveMetastore; -import com.facebook.presto.hive.metastore.MetastoreContext; import com.facebook.presto.hive.metastore.Partition; -import com.facebook.presto.hive.metastore.Table; import com.facebook.presto.hudi.split.ForHudiBackgroundSplitLoader; import com.facebook.presto.hudi.split.ForHudiSplitAsyncQueue; import com.facebook.presto.hudi.split.ForHudiSplitSource; @@ -33,8 +31,6 @@ import com.facebook.presto.spi.connector.ConnectorSplitManager; import com.facebook.presto.spi.connector.ConnectorTransactionHandle; import com.google.common.collect.ImmutableList; -import com.google.common.collect.ImmutableMap; -import com.google.common.collect.Streams; import jakarta.inject.Inject; import org.apache.hadoop.conf.Configuration; import org.apache.hadoop.fs.Path; @@ -45,23 +41,19 @@ import org.apache.hudi.common.table.timeline.HoodieTimeline; import org.apache.hudi.common.table.view.HoodieTableFileSystemView; import org.apache.hudi.common.util.HoodieTimer; +import org.apache.hudi.storage.StorageConfiguration; import java.io.IOException; -import java.util.List; import java.util.Map; import java.util.concurrent.ExecutorService; import java.util.concurrent.ScheduledExecutorService; -import static com.facebook.presto.hive.metastore.MetastoreUtil.extractPartitionValues; import static com.facebook.presto.hudi.HudiErrorCode.HUDI_FILESYSTEM_ERROR; -import static com.facebook.presto.hudi.HudiErrorCode.HUDI_INVALID_METADATA; -import static com.facebook.presto.hudi.HudiMetadata.fromDataColumns; import static com.facebook.presto.hudi.HudiSessionProperties.getMaxOutstandingSplits; import static com.facebook.presto.hudi.HudiSessionProperties.isHudiMetadataTableEnabled; -import static com.google.common.base.Preconditions.checkArgument; -import static java.lang.String.format; import static java.util.Objects.requireNonNull; import static org.apache.hudi.common.table.view.FileSystemViewManager.createInMemoryFileSystemViewWithTimeline; +import static org.apache.hudi.hadoop.fs.HadoopFSUtils.getStorageConfWithCopy; public class HudiSplitManager implements ConnectorSplitManager @@ -104,8 +96,8 @@ public ConnectorSplitSource getSplits( HudiTableHandle table = layout.getTable(); // Retrieve and prune partitions - HoodieTimer timer = new HoodieTimer().startTimer(); - List partitions = hudiPartitionManager.getEffectivePartitions(session, metastore, table.getSchemaTableName(), layout.getTupleDomain()); + HoodieTimer timer = HoodieTimer.start(); + Map partitions = hudiPartitionManager.getEffectivePartitions(session, metastore, table.getSchemaTableName(), table.getPath(), layout.getTupleDomain()); log.debug("Took %d ms to get %d partitions", timer.endTimer(), partitions.size()); if (partitions.isEmpty()) { return new FixedSplitSource(ImmutableList.of()); @@ -114,10 +106,10 @@ public ConnectorSplitSource getSplits( // Load Hudi metadata ExtendedFileSystem fs = getFileSystem(session, table); HoodieMetadataConfig metadataConfig = HoodieMetadataConfig.newBuilder().enable(isHudiMetadataTableEnabled(session)).build(); - Configuration conf = fs.getConf(); + StorageConfiguration conf = getStorageConfWithCopy(fs.getConf()); HoodieTableMetaClient metaClient = HoodieTableMetaClient.builder().setConf(conf).setBasePath(table.getPath()).build(); HoodieTimeline timeline = metaClient.getActiveTimeline().getCommitsTimeline().filterCompletedInstants(); - String timestamp = timeline.lastInstant().map(HoodieInstant::getTimestamp).orElse(null); + String timestamp = timeline.lastInstant().map(HoodieInstant::requestedTime).orElse(null); if (timestamp == null) { // no completed instant for current table return new FixedSplitSource(ImmutableList.of()); @@ -127,7 +119,6 @@ public ConnectorSplitSource getSplits( return new HudiSplitSource( session, - metastore, layout, fsView, partitions, @@ -153,36 +144,4 @@ private ExtendedFileSystem getFileSystem(ConnectorSession session, HudiTableHand throw new PrestoException(HUDI_FILESYSTEM_ERROR, "Could not open file system for " + table, e); } } - - public static HudiPartition getHudiPartition(ExtendedHiveMetastore metastore, MetastoreContext context, HudiTableLayoutHandle tableLayout, String partitionName) - { - String databaseName = tableLayout.getTable().getSchemaName(); - String tableName = tableLayout.getTable().getTableName(); - List partitionColumns = tableLayout.getPartitionColumns(); - - if (partitionColumns.isEmpty()) { - // non-partitioned tableLayout - Table table = metastore.getTable(context, databaseName, tableName) - .orElseThrow(() -> new PrestoException(HUDI_INVALID_METADATA, format("Table %s.%s expected but not found", databaseName, tableName))); - return new HudiPartition(partitionName, ImmutableList.of(), ImmutableMap.of(), table.getStorage(), tableLayout.getDataColumns()); - } - else { - // partitioned tableLayout - List partitionValues = extractPartitionValues(partitionName); - checkArgument(partitionColumns.size() == partitionValues.size(), - format("Invalid partition name %s for partition columns %s", partitionName, partitionColumns)); - Partition partition = metastore.getPartition(context, databaseName, tableName, partitionValues) - .orElseThrow(() -> new PrestoException(HUDI_INVALID_METADATA, format("Partition %s expected but not found", partitionName))); - Map keyValues = zipPartitionKeyValues(partitionColumns, partitionValues); - return new HudiPartition(partitionName, partitionValues, keyValues, partition.getStorage(), fromDataColumns(partition.getColumns())); - } - } - - private static Map zipPartitionKeyValues(List partitionColumns, List partitionValues) - { - ImmutableMap.Builder builder = ImmutableMap.builder(); - Streams.forEachPair(partitionColumns.stream(), partitionValues.stream(), - (column, value) -> builder.put(column.getName(), value)); - return builder.build(); - } } diff --git a/presto-hudi/src/main/java/com/facebook/presto/hudi/HudiSplitSource.java b/presto-hudi/src/main/java/com/facebook/presto/hudi/HudiSplitSource.java index 989870b2bee7b..c69cfa3f3903c 100644 --- a/presto-hudi/src/main/java/com/facebook/presto/hudi/HudiSplitSource.java +++ b/presto-hudi/src/main/java/com/facebook/presto/hudi/HudiSplitSource.java @@ -14,7 +14,7 @@ package com.facebook.presto.hudi; -import com.facebook.presto.hive.metastore.ExtendedHiveMetastore; +import com.facebook.presto.hive.metastore.Partition; import com.facebook.presto.hive.util.AsyncQueue; import com.facebook.presto.hudi.split.HudiBackgroundSplitLoader; import com.facebook.presto.spi.ConnectorSession; @@ -24,7 +24,7 @@ import com.google.common.util.concurrent.Futures; import org.apache.hudi.common.table.view.HoodieTableFileSystemView; -import java.util.List; +import java.util.Map; import java.util.concurrent.CompletableFuture; import java.util.concurrent.ExecutorService; import java.util.concurrent.ScheduledExecutorService; @@ -43,10 +43,9 @@ public class HudiSplitSource public HudiSplitSource( ConnectorSession session, - ExtendedHiveMetastore metastore, HudiTableLayoutHandle layout, HoodieTableFileSystemView fsView, - List partitions, + Map partitions, String latestInstant, ExecutorService asyncQueueExecutor, ScheduledExecutorService splitLoaderExecutorService, @@ -56,7 +55,6 @@ public HudiSplitSource( this.queue = new AsyncQueue<>(maxOutstandingSplits, asyncQueueExecutor); this.splitLoader = new HudiBackgroundSplitLoader( session, - metastore, splitGeneratorExecutorService, layout, fsView, diff --git a/presto-hudi/src/main/java/com/facebook/presto/hudi/HudiTableHandle.java b/presto-hudi/src/main/java/com/facebook/presto/hudi/HudiTableHandle.java index 686267d2d1c8a..d8a21b27f37e3 100644 --- a/presto-hudi/src/main/java/com/facebook/presto/hudi/HudiTableHandle.java +++ b/presto-hudi/src/main/java/com/facebook/presto/hudi/HudiTableHandle.java @@ -14,13 +14,16 @@ package com.facebook.presto.hudi; +import com.facebook.presto.hive.metastore.Table; import com.facebook.presto.spi.ConnectorTableHandle; import com.facebook.presto.spi.SchemaTableName; import com.fasterxml.jackson.annotation.JsonCreator; import com.fasterxml.jackson.annotation.JsonProperty; import java.util.Objects; +import java.util.Optional; +import static com.google.common.base.Preconditions.checkArgument; import static java.util.Objects.requireNonNull; public class HudiTableHandle @@ -30,6 +33,7 @@ public class HudiTableHandle private final String tableName; private final String path; private final HudiTableType hudiTableType; + private final transient Optional
table; @JsonCreator public HudiTableHandle( @@ -38,12 +42,31 @@ public HudiTableHandle( @JsonProperty("path") String path, @JsonProperty("tableType") HudiTableType hudiTableType) { + this(Optional.empty(), schemaName, tableName, path, hudiTableType); + } + + public HudiTableHandle( + Optional
table, + String schemaName, + String tableName, + String path, + HudiTableType hudiTableType) + { + this.table = requireNonNull(table, "table is null"); this.schemaName = requireNonNull(schemaName, "schemaName is null"); this.tableName = requireNonNull(tableName, "tableName is null"); this.path = requireNonNull(path, "path is null"); this.hudiTableType = requireNonNull(hudiTableType, "tableType is null"); } + public Table getTable() + { + checkArgument(table.isPresent(), + "getTable() called on a table handle that has no metastore table object; " + + "this is likely because it is called on the worker."); + return table.get(); + } + @JsonProperty public String getSchemaName() { diff --git a/presto-hudi/src/main/java/com/facebook/presto/hudi/split/HudiBackgroundSplitLoader.java b/presto-hudi/src/main/java/com/facebook/presto/hudi/split/HudiBackgroundSplitLoader.java index 14894adf26318..afbff033d5e6c 100644 --- a/presto-hudi/src/main/java/com/facebook/presto/hudi/split/HudiBackgroundSplitLoader.java +++ b/presto-hudi/src/main/java/com/facebook/presto/hudi/split/HudiBackgroundSplitLoader.java @@ -15,7 +15,7 @@ package com.facebook.presto.hudi.split; import com.facebook.airlift.log.Logger; -import com.facebook.presto.hive.metastore.ExtendedHiveMetastore; +import com.facebook.presto.hive.metastore.Partition; import com.facebook.presto.hive.util.AsyncQueue; import com.facebook.presto.hudi.HudiTableLayoutHandle; import com.facebook.presto.spi.ConnectorSession; @@ -26,6 +26,7 @@ import java.util.ArrayList; import java.util.List; +import java.util.Map; import java.util.concurrent.ConcurrentLinkedQueue; import java.util.concurrent.ExecutionException; import java.util.concurrent.ExecutorService; @@ -44,27 +45,24 @@ public class HudiBackgroundSplitLoader private static final Logger log = Logger.get(HudiBackgroundSplitLoader.class); private final ConnectorSession session; - private final ExtendedHiveMetastore metastore; private final HudiTableLayoutHandle layout; private final HoodieTableFileSystemView fsView; private final AsyncQueue asyncQueue; - private final List partitions; + private final Map partitions; private final String latestInstant; private final int splitGeneratorNumThreads; private final ExecutorService splitGeneratorExecutorService; public HudiBackgroundSplitLoader( ConnectorSession session, - ExtendedHiveMetastore metastore, ExecutorService splitGeneratorExecutorService, HudiTableLayoutHandle layout, HoodieTableFileSystemView fsView, AsyncQueue asyncQueue, - List partitions, + Map partitions, String latestInstant) { this.session = requireNonNull(session, "session is null"); - this.metastore = requireNonNull(metastore, "metastore is null"); this.layout = requireNonNull(layout, "layout is null"); this.fsView = requireNonNull(fsView, "fsView is null"); this.asyncQueue = requireNonNull(asyncQueue, "asyncQueue is null"); @@ -78,15 +76,15 @@ public HudiBackgroundSplitLoader( @Override public void run() { - HoodieTimer timer = new HoodieTimer().startTimer(); + HoodieTimer timer = HoodieTimer.start(); List splitGeneratorList = new ArrayList<>(); List splitGeneratorFutures = new ArrayList<>(); - ConcurrentLinkedQueue concurrentPartitionQueue = new ConcurrentLinkedQueue<>(partitions); + ConcurrentLinkedQueue concurrentPartitionQueue = new ConcurrentLinkedQueue<>(partitions.keySet()); // Start a number of partition split generators to generate the splits in parallel for (int i = 0; i < splitGeneratorNumThreads; i++) { HudiPartitionSplitGenerator generator = new HudiPartitionSplitGenerator( - session, metastore, layout, fsView, asyncQueue, concurrentPartitionQueue, latestInstant); + session, layout, fsView, partitions, asyncQueue, concurrentPartitionQueue, latestInstant); splitGeneratorList.add(generator); splitGeneratorFutures.add(splitGeneratorExecutorService.submit(generator)); } diff --git a/presto-hudi/src/main/java/com/facebook/presto/hudi/split/HudiPartitionSplitGenerator.java b/presto-hudi/src/main/java/com/facebook/presto/hudi/split/HudiPartitionSplitGenerator.java index 54fd0f4bda8a0..d42eae11c84ef 100644 --- a/presto-hudi/src/main/java/com/facebook/presto/hudi/split/HudiPartitionSplitGenerator.java +++ b/presto-hudi/src/main/java/com/facebook/presto/hudi/split/HudiPartitionSplitGenerator.java @@ -16,9 +16,10 @@ import com.facebook.airlift.log.Logger; import com.facebook.airlift.units.DataSize; -import com.facebook.presto.hive.metastore.ExtendedHiveMetastore; -import com.facebook.presto.hive.metastore.MetastoreContext; +import com.facebook.presto.hive.metastore.Partition; +import com.facebook.presto.hive.metastore.Table; import com.facebook.presto.hive.util.AsyncQueue; +import com.facebook.presto.hudi.HudiColumnHandle; import com.facebook.presto.hudi.HudiFile; import com.facebook.presto.hudi.HudiPartition; import com.facebook.presto.hudi.HudiSplit; @@ -27,26 +28,34 @@ import com.facebook.presto.hudi.HudiTableType; import com.facebook.presto.spi.ConnectorSession; import com.facebook.presto.spi.ConnectorSplit; +import com.facebook.presto.spi.PrestoException; import com.facebook.presto.spi.schedule.NodeSelectionStrategy; import com.google.common.collect.ImmutableList; +import com.google.common.collect.ImmutableMap; +import com.google.common.collect.Streams; import org.apache.hadoop.fs.Path; -import org.apache.hudi.common.fs.FSUtils; import org.apache.hudi.common.model.FileSlice; import org.apache.hudi.common.table.view.HoodieTableFileSystemView; import org.apache.hudi.common.util.HoodieTimer; import java.util.List; +import java.util.Map; import java.util.Optional; import java.util.Queue; import java.util.stream.Stream; -import static com.facebook.presto.hudi.HudiMetadata.toMetastoreContext; +import static com.facebook.presto.hive.metastore.MetastoreUtil.extractPartitionValues; +import static com.facebook.presto.hudi.HudiErrorCode.HUDI_INVALID_METADATA; +import static com.facebook.presto.hudi.HudiMetadata.fromDataColumns; import static com.facebook.presto.hudi.HudiSessionProperties.getMinimumAssignedSplitWeight; import static com.facebook.presto.hudi.HudiSessionProperties.getStandardSplitWeightSize; import static com.facebook.presto.hudi.HudiSessionProperties.isSizeBasedSplitWeightsEnabled; -import static com.facebook.presto.hudi.HudiSplitManager.getHudiPartition; +import static com.google.common.base.Preconditions.checkArgument; import static com.google.common.collect.ImmutableList.toImmutableList; +import static java.lang.String.format; import static java.util.Objects.requireNonNull; +import static org.apache.hudi.common.fs.FSUtils.getRelativePartitionPath; +import static org.apache.hudi.hadoop.fs.HadoopFSUtils.convertToStoragePath; /** * A runnable to take partition names from a queue of partitions to process, @@ -58,8 +67,6 @@ public class HudiPartitionSplitGenerator { private static final Logger log = Logger.get(HudiPartitionSplitGenerator.class); - private final ExtendedHiveMetastore metastore; - private final MetastoreContext metastoreContext; private final HudiTableLayoutHandle layout; private final HudiTableHandle table; private final Path tablePath; @@ -68,22 +75,22 @@ public class HudiPartitionSplitGenerator private final Queue concurrentPartitionQueue; private final String latestInstant; private final HudiSplitWeightProvider splitWeightProvider; + private final Map partitions; public HudiPartitionSplitGenerator( ConnectorSession session, - ExtendedHiveMetastore metastore, HudiTableLayoutHandle layout, HoodieTableFileSystemView fsView, + Map partitions, AsyncQueue asyncQueue, Queue concurrentPartitionQueue, String latestInstant) { - this.metastore = requireNonNull(metastore, "metastore is null"); - this.metastoreContext = toMetastoreContext(requireNonNull(session, "session is null")); this.layout = requireNonNull(layout, "layout is null"); this.table = layout.getTable(); this.tablePath = new Path(table.getPath()); this.fsView = requireNonNull(fsView, "fsView is null"); + this.partitions = requireNonNull(partitions, "partitionMap is null"); this.asyncQueue = requireNonNull(asyncQueue, "asyncQueue is null"); this.concurrentPartitionQueue = requireNonNull(concurrentPartitionQueue, "concurrentPartitionQueue is null"); this.latestInstant = requireNonNull(latestInstant, "latestInstant is null"); @@ -93,7 +100,7 @@ public HudiPartitionSplitGenerator( @Override public void run() { - HoodieTimer timer = new HoodieTimer().startTimer(); + HoodieTimer timer = HoodieTimer.start(); while (!concurrentPartitionQueue.isEmpty()) { String partitionName = concurrentPartitionQueue.poll(); if (partitionName != null) { @@ -105,9 +112,9 @@ public void run() private void generateSplitsFromPartition(String partitionName) { - HudiPartition hudiPartition = getHudiPartition(metastore, metastoreContext, layout, partitionName); + HudiPartition hudiPartition = getHudiPartition(layout, partitionName); Path partitionPath = new Path(hudiPartition.getStorage().getLocation()); - String relativePartitionPath = FSUtils.getRelativePartitionPath(tablePath, partitionPath); + String relativePartitionPath = getRelativePartitionPath(convertToStoragePath(tablePath), convertToStoragePath(partitionPath)); Stream fileSlices = HudiTableType.MOR.equals(table.getTableType()) ? fsView.getLatestMergedFileSlicesBeforeOrOn(relativePartitionPath, latestInstant) : fsView.getLatestFileSlicesBeforeOrOn(relativePartitionPath, latestInstant, false); @@ -117,6 +124,38 @@ private void generateSplitsFromPartition(String partitionName) .forEach(asyncQueue::offer); } + private HudiPartition getHudiPartition(HudiTableLayoutHandle tableLayout, String partitionName) + { + String databaseName = tableLayout.getTable().getSchemaName(); + String tableName = tableLayout.getTable().getTableName(); + List partitionColumns = tableLayout.getPartitionColumns(); + + if (partitionColumns.isEmpty()) { + // non-partitioned tableLayout + Table metastoreTable = Optional.ofNullable(table.getTable()) + .orElseThrow(() -> new PrestoException(HUDI_INVALID_METADATA, format("Table %s.%s expected but not found", databaseName, tableName))); + return new HudiPartition(partitionName, ImmutableList.of(), ImmutableMap.of(), metastoreTable.getStorage(), tableLayout.getDataColumns()); + } + else { + // partitioned tableLayout + List partitionValues = extractPartitionValues(partitionName); + checkArgument(partitionColumns.size() == partitionValues.size(), + format("Invalid partition name %s for partition columns %s", partitionName, partitionColumns)); + Partition partition = Optional.ofNullable(partitions.get(partitionName)) + .orElseThrow(() -> new PrestoException(HUDI_INVALID_METADATA, format("Partition %s expected but not found", partitionName))); + Map keyValues = zipPartitionKeyValues(partitionColumns, partitionValues); + return new HudiPartition(partitionName, partitionValues, keyValues, partition.getStorage(), fromDataColumns(partition.getColumns())); + } + } + + private Map zipPartitionKeyValues(List partitionColumns, List partitionValues) + { + ImmutableMap.Builder builder = ImmutableMap.builder(); + Streams.forEachPair(partitionColumns.stream(), partitionValues.stream(), + (column, value) -> builder.put(column.getName(), value)); + return builder.build(); + } + private Optional createHudiSplit( HudiTableHandle table, FileSlice slice, @@ -131,7 +170,7 @@ private Optional createHudiSplit( List logFiles = slice.getLogFiles() .map(logFile -> new HudiFile(logFile.getPath().toString(), 0, logFile.getFileSize())) .collect(toImmutableList()); - long logFilesSize = logFiles.size() > 0 ? logFiles.stream().map(HudiFile::getLength).reduce(0L, Long::sum) : 0L; + long logFilesSize = logFiles.isEmpty() ? 0L : logFiles.stream().mapToLong(HudiFile::getLength).sum(); long sizeInBytes = baseFile != null ? baseFile.getLength() + logFilesSize : logFilesSize; return Optional.of(new HudiSplit( diff --git a/presto-hudi/src/test/java/com/facebook/presto/hudi/TestHudiPartitionManager.java b/presto-hudi/src/test/java/com/facebook/presto/hudi/TestHudiPartitionManager.java index ba445bc1d2054..abe1e3a10dd5a 100644 --- a/presto-hudi/src/test/java/com/facebook/presto/hudi/TestHudiPartitionManager.java +++ b/presto-hudi/src/test/java/com/facebook/presto/hudi/TestHudiPartitionManager.java @@ -23,18 +23,21 @@ import com.facebook.presto.hive.OrcFileWriterConfig; import com.facebook.presto.hive.ParquetFileWriterConfig; import com.facebook.presto.hive.metastore.Column; +import com.facebook.presto.hive.metastore.Partition; import com.facebook.presto.hive.metastore.PrestoTableType; import com.facebook.presto.hive.metastore.Storage; +import com.facebook.presto.hive.metastore.StorageFormat; import com.facebook.presto.hive.metastore.Table; import com.facebook.presto.spi.ColumnHandle; import com.facebook.presto.spi.ConnectorSession; -import com.facebook.presto.spi.SchemaTableName; import com.facebook.presto.testing.TestingConnectorSession; import com.google.common.collect.ImmutableList; import com.google.common.collect.ImmutableMap; +import com.google.common.collect.ImmutableSet; import org.testng.annotations.Test; -import java.util.List; +import java.util.Collections; +import java.util.Map; import java.util.Optional; import static com.facebook.presto.common.type.VarcharType.VARCHAR; @@ -51,6 +54,7 @@ public class TestHudiPartitionManager { private static final String SCHEMA_NAME = "schema"; private static final String TABLE_NAME = "table"; + private static final String NON_PARTITIONED_TABLE_NAME = "non_partitioned_table"; private static final String USER_NAME = "user"; private static final String LOCATION = "somewhere/over/the/rainbow"; private static final Column PARTITION_COLUMN = new Column("ds", HIVE_STRING, Optional.empty(), Optional.empty()); @@ -78,13 +82,65 @@ public class TestHudiPartitionManager Optional.empty(), Optional.empty()); - private static final List PARTITIONS = ImmutableList.of("ds=2019-07-23", "ds=2019-08-23"); + private static final Map> PARTITION_MAP = ImmutableMap.of( + "ds=2019-07-23", + Optional.of( + Partition.builder() + .setCatalogName(Optional.empty()) + .setDatabaseName(SCHEMA_NAME) + .setTableName(TABLE_NAME) + .withStorage(storageBuilder -> + storageBuilder.setLocation(LOCATION) + .setStorageFormat(StorageFormat.VIEW_STORAGE_FORMAT)) + .setColumns(ImmutableList.of(PARTITION_COLUMN)) + .setValues(Collections.singletonList("2019-07-23")) + .build()), + "ds=2019-08-23", + Optional.of( + Partition.builder() + .setCatalogName(Optional.empty()) + .setDatabaseName(SCHEMA_NAME) + .setTableName(TABLE_NAME) + .withStorage(storageBuilder -> + storageBuilder.setLocation(LOCATION) + .setStorageFormat(StorageFormat.VIEW_STORAGE_FORMAT)) + .setColumns(ImmutableList.of(PARTITION_COLUMN)) + .setValues(Collections.singletonList("2019-08-23")) + .build())); + + private static final Table NON_PARTITIONED_TABLE = new Table( + Optional.of("catalogName"), + SCHEMA_NAME, + NON_PARTITIONED_TABLE_NAME, + USER_NAME, + PrestoTableType.MANAGED_TABLE, + new Storage(fromHiveStorageFormat(PARQUET), + LOCATION, + Optional.of(new HiveBucketProperty( + ImmutableList.of(BUCKET_COLUMN.getName()), + 100, + ImmutableList.of(), + HIVE_COMPATIBLE, + Optional.empty())), + false, + ImmutableMap.of(), + ImmutableMap.of()), + ImmutableList.of(BUCKET_COLUMN), + ImmutableList.of(), + ImmutableMap.of(), + Optional.empty(), + Optional.empty()); + + private static final Map> NON_PARTITION_MAP = ImmutableMap.of( + "", + Optional.empty()); + private final HudiPartitionManager hudiPartitionManager = new HudiPartitionManager(new TestingTypeManager()); - private final TestingExtendedHiveMetastore metastore = new TestingExtendedHiveMetastore(TABLE, PARTITIONS); @Test public void testParseValuesAndFilterPartition() { + TestingExtendedHiveMetastore metastore = new TestingExtendedHiveMetastore(TABLE, PARTITION_MAP); ConnectorSession session = new TestingConnectorSession( new HiveSessionProperties( new HiveClientConfig().setMaxBucketsForGroupedExecution(100), @@ -100,11 +156,37 @@ public void testParseValuesAndFilterPartition() Optional.empty(), HudiColumnHandle.ColumnType.PARTITION_KEY), Domain.singleValue(VARCHAR, utf8Slice("2019-07-23")))); - List actualPartitions = hudiPartitionManager.getEffectivePartitions( + HudiTableHandle tableHandle = new HudiTableHandle(SCHEMA_NAME, TABLE_NAME, LOCATION, HudiTableType.COW); + Map actualPartitions = hudiPartitionManager.getEffectivePartitions( + session, + metastore, + tableHandle.getSchemaTableName(), + tableHandle.getPath(), + constraintSummary); + assertEquals(actualPartitions.keySet(), ImmutableSet.of("ds=2019-07-23")); + assertEquals(actualPartitions.get("ds=2019-07-23"), PARTITION_MAP.get("ds=2019-07-23").get()); + + // Non-partitioned table case + HudiTableHandle nonPartitionedTableHandle = new HudiTableHandle(SCHEMA_NAME, NON_PARTITIONED_TABLE_NAME, LOCATION, HudiTableType.COW); + metastore = new TestingExtendedHiveMetastore(NON_PARTITIONED_TABLE, NON_PARTITION_MAP); + actualPartitions = hudiPartitionManager.getEffectivePartitions( session, metastore, - new SchemaTableName(SCHEMA_NAME, TABLE_NAME), + nonPartitionedTableHandle.getSchemaTableName(), + nonPartitionedTableHandle.getPath(), constraintSummary); - assertEquals(actualPartitions, ImmutableList.of("ds=2019-07-23")); + assertEquals(actualPartitions.keySet(), ImmutableSet.of("")); + assertEquals( + actualPartitions.get(""), + Partition.builder() + .setCatalogName(Optional.empty()) + .setDatabaseName(nonPartitionedTableHandle.getSchemaName()) + .setTableName(nonPartitionedTableHandle.getTableName()) + .withStorage(storageBuilder -> + storageBuilder.setLocation(LOCATION) + .setStorageFormat(StorageFormat.VIEW_STORAGE_FORMAT)) + .setColumns(ImmutableList.of()) + .setValues(ImmutableList.of()) + .build()); } } diff --git a/presto-hudi/src/test/java/com/facebook/presto/hudi/TestingExtendedHiveMetastore.java b/presto-hudi/src/test/java/com/facebook/presto/hudi/TestingExtendedHiveMetastore.java index 40f7fe9425475..ac34d1d70d37f 100644 --- a/presto-hudi/src/test/java/com/facebook/presto/hudi/TestingExtendedHiveMetastore.java +++ b/presto-hudi/src/test/java/com/facebook/presto/hudi/TestingExtendedHiveMetastore.java @@ -18,6 +18,7 @@ import com.facebook.presto.hive.PartitionNameWithVersion; import com.facebook.presto.hive.metastore.Column; import com.facebook.presto.hive.metastore.MetastoreContext; +import com.facebook.presto.hive.metastore.Partition; import com.facebook.presto.hive.metastore.Table; import com.facebook.presto.hive.metastore.UnimplementedHiveMetastore; @@ -26,15 +27,16 @@ import java.util.Optional; import static com.facebook.presto.hive.metastore.MetastoreUtil.getPartitionNamesWithEmptyVersion; +import static com.google.common.collect.ImmutableMap.toImmutableMap; import static java.util.Objects.requireNonNull; public class TestingExtendedHiveMetastore extends UnimplementedHiveMetastore { private final Table table; - private final List partitions; + private final Map> partitions; - public TestingExtendedHiveMetastore(Table table, List partitions) + public TestingExtendedHiveMetastore(Table table, Map> partitions) { this.table = requireNonNull(table, "table is null"); this.partitions = requireNonNull(partitions, "partitions is null"); @@ -53,6 +55,13 @@ public List getPartitionNamesByFilter( String tableName, Map partitionPredicates) { - return getPartitionNamesWithEmptyVersion(partitions); + return getPartitionNamesWithEmptyVersion(partitions.keySet()); + } + + @Override + public Map> getPartitionsByNames(MetastoreContext metastoreContext, String databaseName, String tableName, List partitionNames) + { + return partitionNames.stream().filter(partitionName -> partitions.containsKey(partitionName.getPartitionName())) + .collect(toImmutableMap(PartitionNameWithVersion::getPartitionName, partitionName -> partitions.get(partitionName.getPartitionName()))); } } diff --git a/presto-hudi/src/test/java/com/facebook/presto/hudi/TestingTypeManager.java b/presto-hudi/src/test/java/com/facebook/presto/hudi/TestingTypeManager.java index 8b9f9fb7a747a..c4a0658622fc6 100644 --- a/presto-hudi/src/test/java/com/facebook/presto/hudi/TestingTypeManager.java +++ b/presto-hudi/src/test/java/com/facebook/presto/hudi/TestingTypeManager.java @@ -63,4 +63,10 @@ public List getTypes() { return ImmutableList.of(BOOLEAN, INTEGER, BIGINT, DOUBLE, VARCHAR, VARBINARY, TIMESTAMP, DATE, HYPER_LOG_LOG); } + + @Override + public boolean hasType(TypeSignature signature) + { + return getType(signature) != null; + } } diff --git a/presto-i18n-functions/src/test/java/com/facebook/presto/i18n/functions/TestMyanmarFunctions.java b/presto-i18n-functions/src/test/java/com/facebook/presto/i18n/functions/TestMyanmarFunctions.java index 1f5bc00908408..f0259c46f8c5c 100644 --- a/presto-i18n-functions/src/test/java/com/facebook/presto/i18n/functions/TestMyanmarFunctions.java +++ b/presto-i18n-functions/src/test/java/com/facebook/presto/i18n/functions/TestMyanmarFunctions.java @@ -14,15 +14,23 @@ package com.facebook.presto.i18n.functions; import com.facebook.presto.operator.scalar.AbstractTestFunctions; +import com.facebook.presto.sql.analyzer.FeaturesConfig; +import com.facebook.presto.sql.analyzer.FunctionsConfig; import org.testng.annotations.BeforeClass; import org.testng.annotations.Test; +import static com.facebook.presto.SessionTestUtils.TEST_SESSION; import static com.facebook.presto.common.type.VarcharType.VARCHAR; import static com.facebook.presto.metadata.FunctionExtractor.extractFunctions; public class TestMyanmarFunctions extends AbstractTestFunctions { + public TestMyanmarFunctions() + { + super(TEST_SESSION, new FeaturesConfig(), new FunctionsConfig(), false); + } + @BeforeClass public void setUp() { diff --git a/presto-iceberg/pom.xml b/presto-iceberg/pom.xml index 66445b698a25c..59a435bbff114 100644 --- a/presto-iceberg/pom.xml +++ b/presto-iceberg/pom.xml @@ -14,8 +14,8 @@ ${project.parent.basedir} 17 - 1.6.1 - 0.95.0 + 1.8.1 + 0.103.0 true @@ -613,7 +613,7 @@ org.apache.httpcomponents.core5 httpcore5 - 5.2.4 + 5.3.1 diff --git a/presto-iceberg/src/main/java/com/facebook/presto/iceberg/HiveTableOperations.java b/presto-iceberg/src/main/java/com/facebook/presto/iceberg/HiveTableOperations.java index b608c5e93d332..29b85a11e4fc1 100644 --- a/presto-iceberg/src/main/java/com/facebook/presto/iceberg/HiveTableOperations.java +++ b/presto-iceberg/src/main/java/com/facebook/presto/iceberg/HiveTableOperations.java @@ -30,15 +30,18 @@ import com.facebook.presto.spi.SchemaTableName; import com.facebook.presto.spi.TableNotFoundException; import com.facebook.presto.spi.security.PrestoPrincipal; +import com.google.common.annotations.VisibleForTesting; import com.google.common.cache.CacheBuilder; import com.google.common.cache.CacheLoader; import com.google.common.cache.LoadingCache; +import com.google.common.collect.ImmutableMap; import com.google.common.collect.ImmutableMultimap; import com.google.common.collect.Sets; import jakarta.annotation.Nullable; import org.apache.hadoop.hive.serde2.lazy.LazySimpleSerDe; import org.apache.hadoop.mapred.FileInputFormat; import org.apache.hadoop.mapred.FileOutputFormat; +import org.apache.iceberg.BaseMetastoreTableOperations; import org.apache.iceberg.LocationProviders; import org.apache.iceberg.TableMetadata; import org.apache.iceberg.TableMetadata.MetadataLogEntry; @@ -106,6 +109,7 @@ public class HiveTableOperations private final Optional owner; private final Optional location; private final HdfsFileIO fileIO; + private final IcebergHiveTableOperationsConfig config; private TableMetadata currentMetadata; @@ -251,14 +255,19 @@ public void commit(@Nullable TableMetadata base, TableMetadata metadata) String newMetadataLocation = writeNewMetadata(metadata, version + 1); Table table; - // getting a process-level lock per table to avoid concurrent commit attempts to the same table from the same - // JVM process, which would result in unnecessary and costly HMS lock acquisition requests Optional lockId = Optional.empty(); + boolean useHMSLock = Optional.ofNullable(metadata.property(TableProperties.HIVE_LOCK_ENABLED, null)) + .map(Boolean::parseBoolean) + .orElse(config.getLockingEnabled()); ReentrantLock tableLevelMutex = commitLockCache.getUnchecked(database + "." + tableName); + // getting a process-level lock per table to avoid concurrent commit attempts to the same table from the same + // JVM process, which would result in unnecessary and costly HMS lock acquisition requests tableLevelMutex.lock(); try { try { - lockId = metastore.lock(metastoreContext, database, tableName); + if (useHMSLock) { + lockId = metastore.lock(metastoreContext, database, tableName); + } if (base == null) { String tableComment = metadata.properties().get(TABLE_COMMENT); Map parameters = new HashMap<>(); @@ -318,10 +327,7 @@ public void commit(@Nullable TableMetadata base, TableMetadata metadata) } else { PartitionStatistics tableStats = metastore.getTableStatistics(metastoreContext, database, tableName); - metastore.replaceTable(metastoreContext, database, tableName, table, privileges); - - // attempt to put back previous table statistics - metastore.updateTableStatistics(metastoreContext, database, tableName, oldStats -> tableStats); + metastore.persistTable(metastoreContext, database, tableName, table, privileges, () -> tableStats, useHMSLock ? ImmutableMap.of() : hmsEnvContext(base.metadataFileLocation())); } deleteRemovedMetadataFiles(base, metadata); } @@ -500,4 +506,19 @@ private void deleteRemovedMetadataFiles(TableMetadata base, TableMetadata metada .run(previousMetadataFile -> io().deleteFile(previousMetadataFile.file())); } } + + private Map hmsEnvContext(String metadataLocation) + { + return ImmutableMap.of( + org.apache.iceberg.hive.HiveTableOperations.NO_LOCK_EXPECTED_KEY, + BaseMetastoreTableOperations.METADATA_LOCATION_PROP, + org.apache.iceberg.hive.HiveTableOperations.NO_LOCK_EXPECTED_VALUE, + metadataLocation); + } + + @VisibleForTesting + public IcebergHiveTableOperationsConfig getConfig() + { + return config; + } } diff --git a/presto-iceberg/src/main/java/com/facebook/presto/iceberg/IcebergConnector.java b/presto-iceberg/src/main/java/com/facebook/presto/iceberg/IcebergConnector.java index 6936b7fee96be..b48f916db6e6b 100644 --- a/presto-iceberg/src/main/java/com/facebook/presto/iceberg/IcebergConnector.java +++ b/presto-iceberg/src/main/java/com/facebook/presto/iceberg/IcebergConnector.java @@ -15,6 +15,7 @@ import com.facebook.airlift.bootstrap.LifeCycleManager; import com.facebook.presto.hive.HiveTransactionHandle; +import com.facebook.presto.iceberg.function.IcebergBucketFunction; import com.facebook.presto.iceberg.function.changelog.ApplyChangelogFunction; import com.facebook.presto.spi.SystemTable; import com.facebook.presto.spi.classloader.ThreadContextClassLoader; @@ -225,6 +226,8 @@ public Set> getSystemFunctions() { return ImmutableSet.>builder() .add(ApplyChangelogFunction.class) + .add(IcebergBucketFunction.class) + .add(IcebergBucketFunction.Bucket.class) .build(); } } diff --git a/presto-iceberg/src/main/java/com/facebook/presto/iceberg/IcebergHiveTableOperationsConfig.java b/presto-iceberg/src/main/java/com/facebook/presto/iceberg/IcebergHiveTableOperationsConfig.java index a0cb74928b530..9d9d358044030 100644 --- a/presto-iceberg/src/main/java/com/facebook/presto/iceberg/IcebergHiveTableOperationsConfig.java +++ b/presto-iceberg/src/main/java/com/facebook/presto/iceberg/IcebergHiveTableOperationsConfig.java @@ -31,6 +31,7 @@ public class IcebergHiveTableOperationsConfig private Duration tableRefreshMaxRetryTime = succinctDuration(1, MINUTES); private double tableRefreshBackoffScaleFactor = 4.0; private int tableRefreshRetries = 20; + private boolean lockingEnabled = true; @MinDuration("1ms") public Duration getTableRefreshBackoffMinSleepTime() @@ -101,4 +102,17 @@ public int getTableRefreshRetries() { return tableRefreshRetries; } + + @Config("iceberg.engine.hive.lock-enabled") + @ConfigDescription("Whether to use HMS locks to ensure atomicity of commits") + public IcebergHiveTableOperationsConfig setLockingEnabled(boolean lockingEnabled) + { + this.lockingEnabled = lockingEnabled; + return this; + } + + public boolean getLockingEnabled() + { + return lockingEnabled; + } } diff --git a/presto-iceberg/src/main/java/com/facebook/presto/iceberg/function/IcebergBucketFunction.java b/presto-iceberg/src/main/java/com/facebook/presto/iceberg/function/IcebergBucketFunction.java new file mode 100644 index 0000000000000..429b296f79d8e --- /dev/null +++ b/presto-iceberg/src/main/java/com/facebook/presto/iceberg/function/IcebergBucketFunction.java @@ -0,0 +1,111 @@ +/* + * Licensed 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 com.facebook.presto.iceberg.function; + +import com.facebook.presto.common.type.StandardTypes; +import com.facebook.presto.spi.function.LiteralParameter; +import com.facebook.presto.spi.function.LiteralParameters; +import com.facebook.presto.spi.function.ScalarFunction; +import com.facebook.presto.spi.function.SqlType; +import io.airlift.slice.Slice; +import org.apache.iceberg.transforms.Transforms; +import org.apache.iceberg.types.Types; + +import java.math.BigDecimal; +import java.math.MathContext; + +import static com.facebook.presto.common.type.DateTimeEncoding.unpackMillisUtc; +import static com.facebook.presto.common.type.Decimals.decodeUnscaledValue; +import static com.facebook.presto.common.type.SqlTimestamp.MICROSECONDS_PER_MILLISECOND; + +public final class IcebergBucketFunction +{ + private IcebergBucketFunction() {} + + @ScalarFunction("bucket") + @SqlType(StandardTypes.BIGINT) + public static long bucketInteger(@SqlType(StandardTypes.BIGINT) long value, @SqlType(StandardTypes.INTEGER) long numberOfBuckets) + { + return Transforms.bucket((int) numberOfBuckets) + .bind(Types.LongType.get()) + .apply(value); + } + + @ScalarFunction("bucket") + @SqlType(StandardTypes.BIGINT) + public static long bucketVarchar(@SqlType(StandardTypes.VARCHAR) Slice value, @SqlType(StandardTypes.INTEGER) long numberOfBuckets) + { + return (long) Transforms.bucket((int) numberOfBuckets) + .bind(Types.StringType.get()) + .apply(value.toStringUtf8()); + } + + @ScalarFunction("bucket") + @SqlType(StandardTypes.BIGINT) + public static long bucketVarbinary(@SqlType(StandardTypes.VARBINARY) Slice value, @SqlType(StandardTypes.INTEGER) long numberOfBuckets) + { + return (long) Transforms.bucket((int) numberOfBuckets) + .bind(Types.BinaryType.get()) + .apply(value.toByteBuffer()); + } + + @ScalarFunction("bucket") + @SqlType(StandardTypes.BIGINT) + public static long bucketDate(@SqlType(StandardTypes.DATE) long value, @SqlType(StandardTypes.INTEGER) long numberOfBuckets) + { + return Transforms.bucket((int) numberOfBuckets) + .bind(Types.DateType.get()) + .apply((int) value); + } + + @ScalarFunction("bucket") + @SqlType(StandardTypes.BIGINT) + public static long bucketTimestamp(@SqlType(StandardTypes.TIMESTAMP) long value, @SqlType(StandardTypes.INTEGER) long numberOfBuckets) + { + return Transforms.bucket((int) numberOfBuckets) + .bind(Types.TimestampType.withoutZone()) + .apply(value); + } + + @ScalarFunction("bucket") + @SqlType(StandardTypes.BIGINT) + public static long bucketTimestampWithTimeZone(@SqlType(StandardTypes.TIMESTAMP_WITH_TIME_ZONE) long value, @SqlType(StandardTypes.INTEGER) long numberOfBuckets) + { + return Transforms.bucket((int) numberOfBuckets) + .bind(Types.TimestampType.withZone()) + .apply(unpackMillisUtc(value) * MICROSECONDS_PER_MILLISECOND); + } + + @ScalarFunction("bucket") + public static final class Bucket + { + @LiteralParameters({"p", "s"}) + @SqlType(StandardTypes.BIGINT) + public static long bucketShortDecimal(@LiteralParameter("p") long numPrecision, @LiteralParameter("s") long numScale, @SqlType("decimal(p, s)") long value, @SqlType(StandardTypes.INTEGER) long numberOfBuckets) + { + return Transforms.bucket((int) numberOfBuckets) + .bind(Types.DecimalType.of((int) numPrecision, (int) numScale)) + .apply(BigDecimal.valueOf(value)); + } + + @LiteralParameters({"p", "s"}) + @SqlType(StandardTypes.BIGINT) + public static long bucketLongDecimal(@LiteralParameter("p") long numPrecision, @LiteralParameter("s") long numScale, @SqlType("decimal(p, s)") Slice value, @SqlType(StandardTypes.INTEGER) long numberOfBuckets) + { + return Transforms.bucket((int) numberOfBuckets) + .bind(Types.DecimalType.of((int) numPrecision, (int) numScale)) + .apply(new BigDecimal(decodeUnscaledValue(value), (int) numScale, new MathContext((int) numPrecision))); + } + } +} diff --git a/presto-iceberg/src/test/java/com/facebook/presto/iceberg/TestIcebergDistributedQueries.java b/presto-iceberg/src/test/java/com/facebook/presto/iceberg/TestIcebergDistributedQueries.java index c689545433980..57dd62d142e3a 100644 --- a/presto-iceberg/src/test/java/com/facebook/presto/iceberg/TestIcebergDistributedQueries.java +++ b/presto-iceberg/src/test/java/com/facebook/presto/iceberg/TestIcebergDistributedQueries.java @@ -115,6 +115,18 @@ public void testDescribeOutputNamedAndUnnamed() assertEqualsIgnoreOrder(actual, expected); } + @Override + public void testNonAutoCommitTransactionWithRollback() + { + // Catalog iceberg only supports writes using autocommit + } + + @Override + public void testNonAutoCommitTransactionWithCommit() + { + // Catalog iceberg only supports writes using autocommit + } + /** * Increased the optimizer timeout from 15000ms to 25000ms */ diff --git a/presto-iceberg/src/test/java/com/facebook/presto/iceberg/TestIcebergFileWriter.java b/presto-iceberg/src/test/java/com/facebook/presto/iceberg/TestIcebergFileWriter.java index 4e14cf3a1c639..31f0729907da2 100644 --- a/presto-iceberg/src/test/java/com/facebook/presto/iceberg/TestIcebergFileWriter.java +++ b/presto-iceberg/src/test/java/com/facebook/presto/iceberg/TestIcebergFileWriter.java @@ -189,5 +189,11 @@ public List getTypes() { return ImmutableList.of(BooleanType.BOOLEAN, INTEGER, BIGINT, DoubleType.DOUBLE, VARCHAR, VARBINARY, TIMESTAMP, DATE, HYPER_LOG_LOG); } + + @Override + public boolean hasType(TypeSignature signature) + { + return getType(signature) != null; + } } } diff --git a/presto-iceberg/src/test/java/com/facebook/presto/iceberg/TestIcebergScalarFunctions.java b/presto-iceberg/src/test/java/com/facebook/presto/iceberg/TestIcebergScalarFunctions.java new file mode 100644 index 0000000000000..c6253afcb9638 --- /dev/null +++ b/presto-iceberg/src/test/java/com/facebook/presto/iceberg/TestIcebergScalarFunctions.java @@ -0,0 +1,80 @@ +/* + * Licensed 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 com.facebook.presto.iceberg; + +import com.facebook.presto.common.CatalogSchemaName; +import com.facebook.presto.iceberg.function.IcebergBucketFunction; +import com.facebook.presto.metadata.FunctionExtractor; +import com.facebook.presto.operator.scalar.AbstractTestFunctions; +import com.facebook.presto.sql.analyzer.FeaturesConfig; +import com.facebook.presto.sql.analyzer.FunctionsConfig; +import com.facebook.presto.type.DateOperators; +import com.facebook.presto.type.TimestampOperators; +import com.facebook.presto.type.TimestampWithTimeZoneOperators; +import org.testcontainers.shaded.com.google.common.collect.ImmutableList; +import org.testng.annotations.BeforeClass; +import org.testng.annotations.Test; + +import java.math.BigDecimal; + +import static com.facebook.presto.SessionTestUtils.TEST_SESSION; +import static com.facebook.presto.common.type.BigintType.BIGINT; +import static com.facebook.presto.common.type.Decimals.encodeScaledValue; +import static com.facebook.presto.iceberg.function.IcebergBucketFunction.Bucket.bucketLongDecimal; +import static com.facebook.presto.iceberg.function.IcebergBucketFunction.Bucket.bucketShortDecimal; +import static com.facebook.presto.iceberg.function.IcebergBucketFunction.bucketDate; +import static com.facebook.presto.iceberg.function.IcebergBucketFunction.bucketInteger; +import static com.facebook.presto.iceberg.function.IcebergBucketFunction.bucketTimestamp; +import static com.facebook.presto.iceberg.function.IcebergBucketFunction.bucketTimestampWithTimeZone; +import static com.facebook.presto.iceberg.function.IcebergBucketFunction.bucketVarbinary; +import static com.facebook.presto.iceberg.function.IcebergBucketFunction.bucketVarchar; +import static io.airlift.slice.Slices.utf8Slice; + +public class TestIcebergScalarFunctions + extends AbstractTestFunctions +{ + public TestIcebergScalarFunctions() + { + super(TEST_SESSION, new FeaturesConfig(), new FunctionsConfig(), false); + } + + @BeforeClass + public void registerFunction() + { + ImmutableList.Builder> functions = ImmutableList.builder(); + functions.add(IcebergBucketFunction.class) + .add(IcebergBucketFunction.Bucket.class); + functionAssertions.addConnectorFunctions(FunctionExtractor.extractFunctions(functions.build(), + new CatalogSchemaName("iceberg", "system")), "iceberg"); + } + + @Test + public void testBucketFunction() + { + String catalogSchema = "iceberg.system"; + functionAssertions.assertFunction(catalogSchema + ".bucket(cast(10 as tinyint), 3)", BIGINT, bucketInteger(10, 3)); + functionAssertions.assertFunction(catalogSchema + ".bucket(cast(1950 as smallint), 4)", BIGINT, bucketInteger(1950, 4)); + functionAssertions.assertFunction(catalogSchema + ".bucket(cast(2375645 as int), 5)", BIGINT, bucketInteger(2375645, 5)); + functionAssertions.assertFunction(catalogSchema + ".bucket(cast(2779099983928392323 as bigint), 6)", BIGINT, bucketInteger(2779099983928392323L, 6)); + functionAssertions.assertFunction(catalogSchema + ".bucket(cast(456.43 as DECIMAL(5,2)), 12)", BIGINT, bucketShortDecimal(5, 2, 45643, 12)); + functionAssertions.assertFunction(catalogSchema + ".bucket(cast('12345678901234567890.1234567890' as DECIMAL(30,10)), 12)", BIGINT, bucketLongDecimal(30, 10, encodeScaledValue(new BigDecimal("12345678901234567890.1234567890")), 12)); + + functionAssertions.assertFunction(catalogSchema + ".bucket(cast('nasdbsdnsdms' as varchar), 7)", BIGINT, bucketVarchar(utf8Slice("nasdbsdnsdms"), 7)); + functionAssertions.assertFunction(catalogSchema + ".bucket(cast('nasdbsdnsdms' as varbinary), 8)", BIGINT, bucketVarbinary(utf8Slice("nasdbsdnsdms"), 8)); + + functionAssertions.assertFunction(catalogSchema + ".bucket(cast('2018-04-06' as date), 9)", BIGINT, bucketDate(DateOperators.castFromSlice(utf8Slice("2018-04-06")), 9)); + functionAssertions.assertFunction(catalogSchema + ".bucket(CAST('2018-04-06 04:35:00.000' AS TIMESTAMP),10)", BIGINT, bucketTimestamp(TimestampOperators.castFromSlice(TEST_SESSION.getSqlFunctionProperties(), utf8Slice("2018-04-06 04:35:00.000")), 10)); + functionAssertions.assertFunction(catalogSchema + ".bucket(CAST('2018-04-06 04:35:00.000 GMT' AS TIMESTAMP WITH TIME ZONE), 11)", BIGINT, bucketTimestampWithTimeZone(TimestampWithTimeZoneOperators.castFromSlice(TEST_SESSION.getSqlFunctionProperties(), utf8Slice("2018-04-06 04:35:00.000 GMT")), 11)); + } +} diff --git a/presto-iceberg/src/test/java/com/facebook/presto/iceberg/TestOutputColumnTypes.java b/presto-iceberg/src/test/java/com/facebook/presto/iceberg/TestOutputColumnTypes.java index 8f247835e7418..89551f5d388a7 100644 --- a/presto-iceberg/src/test/java/com/facebook/presto/iceberg/TestOutputColumnTypes.java +++ b/presto-iceberg/src/test/java/com/facebook/presto/iceberg/TestOutputColumnTypes.java @@ -15,10 +15,12 @@ package com.facebook.presto.iceberg; import com.facebook.presto.Session; +import com.facebook.presto.common.QualifiedObjectName; +import com.facebook.presto.common.SourceColumn; import com.facebook.presto.spi.Plugin; -import com.facebook.presto.spi.eventlistener.Column; import com.facebook.presto.spi.eventlistener.EventListener; import com.facebook.presto.spi.eventlistener.EventListenerFactory; +import com.facebook.presto.spi.eventlistener.OutputColumnMetadata; import com.facebook.presto.spi.eventlistener.QueryCompletedEvent; import com.facebook.presto.spi.eventlistener.QueryCreatedEvent; import com.facebook.presto.spi.eventlistener.SplitCompletedEvent; @@ -26,6 +28,7 @@ import com.facebook.presto.testing.QueryRunner; import com.facebook.presto.tests.AbstractTestQueryFramework; import com.google.common.collect.ImmutableList; +import com.google.common.collect.ImmutableSet; import org.intellij.lang.annotations.Language; import org.testng.annotations.AfterClass; import org.testng.annotations.Test; @@ -108,34 +111,214 @@ public void testOutputColumnsForInsertAsSelect() assertThat(event.getIoMetadata().getOutput().get().getColumns().get()) .containsExactly( - new Column("clerk", "varchar"), - new Column("orderkey", "bigint"), - new Column("totalprice", "double")); + new OutputColumnMetadata("clerk", "varchar", ImmutableSet.of( + new SourceColumn(new QualifiedObjectName("iceberg", "tpch", "orders"), "clerk"))), + new OutputColumnMetadata("orderkey", "bigint", ImmutableSet.of( + new SourceColumn(new QualifiedObjectName("iceberg", "tpch", "orders"), "orderkey"))), + new OutputColumnMetadata("totalprice", "double", ImmutableSet.of( + new SourceColumn(new QualifiedObjectName("iceberg", "tpch", "orders"), "totalprice")))); + } + + @Test + public void testOutputColumnsForInsertAsSelectAllWithAliasedRelation() + throws Exception + { + runQueryAndWaitForEvents("CREATE TABLE create_insert_output1 AS SELECT clerk AS test_clerk, orderkey AS test_orderkey, totalprice AS test_totalprice FROM orders", 2); + runQueryAndWaitForEvents("INSERT INTO create_insert_output1(test_clerk,test_orderkey,test_totalprice) SELECT clerk AS test_clerk, orderkey AS test_orderkey, totalprice AS test_totalprice FROM (SELECT * from orders) orders_a", 2); + QueryCompletedEvent event = getOnlyElement(generatedEvents.getQueryCompletedEvents()); + + assertThat(event.getIoMetadata().getOutput().get().getCatalogName()).isEqualTo("iceberg"); + assertThat(event.getIoMetadata().getOutput().get().getSchema()).isEqualTo("tpch"); + assertThat(event.getIoMetadata().getOutput().get().getTable()).isEqualTo("create_insert_output1"); + assertThat(event.getMetadata().getUpdateQueryType().get()).isEqualTo("INSERT"); + + assertThat(event.getIoMetadata().getOutput().get().getColumns().get()) + .containsExactly( + new OutputColumnMetadata("test_clerk", "varchar", ImmutableSet.of( + new SourceColumn(new QualifiedObjectName("iceberg", "tpch", "orders"), "clerk"))), + new OutputColumnMetadata("test_orderkey", "bigint", ImmutableSet.of( + new SourceColumn(new QualifiedObjectName("iceberg", "tpch", "orders"), "orderkey"))), + new OutputColumnMetadata("test_totalprice", "double", ImmutableSet.of( + new SourceColumn(new QualifiedObjectName("iceberg", "tpch", "orders"), "totalprice")))); + } + + @Test + public void testOutputColumnsForInsertAsSelectColumnAliasInAliasedRelation() + throws Exception + { + runQueryAndWaitForEvents("CREATE TABLE create_insert_output2 AS SELECT clerk AS test_clerk, orderkey AS test_orderkey, totalprice AS test_totalprice FROM orders", 2); + runQueryAndWaitForEvents("INSERT INTO create_insert_output2(test_clerk,test_orderkey,test_totalprice) SELECT aliased_clerk AS test_clerk, aliased_orderkey AS test_orderkey, aliased_totalprice AS test_totalprice FROM (SELECT clerk, orderkey, totalprice from orders) orders_a(aliased_clerk, aliased_orderkey, aliased_totalprice)", 2); + QueryCompletedEvent event = getOnlyElement(generatedEvents.getQueryCompletedEvents()); + + assertThat(event.getIoMetadata().getOutput().get().getCatalogName()).isEqualTo("iceberg"); + assertThat(event.getIoMetadata().getOutput().get().getSchema()).isEqualTo("tpch"); + assertThat(event.getIoMetadata().getOutput().get().getTable()).isEqualTo("create_insert_output2"); + assertThat(event.getMetadata().getUpdateQueryType().get()).isEqualTo("INSERT"); + + assertThat(event.getIoMetadata().getOutput().get().getColumns().get()) + .containsExactly( + new OutputColumnMetadata("test_clerk", "varchar", ImmutableSet.of( + new SourceColumn(new QualifiedObjectName("iceberg", "tpch", "orders"), "clerk"))), + new OutputColumnMetadata("test_orderkey", "bigint", ImmutableSet.of( + new SourceColumn(new QualifiedObjectName("iceberg", "tpch", "orders"), "orderkey"))), + new OutputColumnMetadata("test_totalprice", "double", ImmutableSet.of( + new SourceColumn(new QualifiedObjectName("iceberg", "tpch", "orders"), "totalprice")))); } @Test public void testOutputColumnsForCreateTableAS() throws Exception { - runQueryAndWaitForEvents("CREATE TABLE create_update_table AS SELECT * FROM orders ", 2); + runQueryAndWaitForEvents("CREATE TABLE create_update_table2 AS SELECT * FROM orders ", 2); + QueryCompletedEvent event = getOnlyElement(generatedEvents.getQueryCompletedEvents()); + + assertThat(event.getIoMetadata().getOutput().get().getCatalogName()).isEqualTo("iceberg"); + assertThat(event.getIoMetadata().getOutput().get().getSchema()).isEqualTo("tpch"); + assertThat(event.getIoMetadata().getOutput().get().getTable()).isEqualTo("create_update_table2"); + assertThat(event.getMetadata().getUpdateQueryType().get()).isEqualTo("CREATE TABLE"); + + assertThat(event.getIoMetadata().getOutput().get().getColumns().get()) + .containsExactly( + new OutputColumnMetadata("orderkey", "bigint", ImmutableSet.of( + new SourceColumn(new QualifiedObjectName("iceberg", "tpch", "orders"), "orderkey"))), + new OutputColumnMetadata("custkey", "bigint", ImmutableSet.of( + new SourceColumn(new QualifiedObjectName("iceberg", "tpch", "orders"), "custkey"))), + new OutputColumnMetadata("orderstatus", "varchar", ImmutableSet.of( + new SourceColumn(new QualifiedObjectName("iceberg", "tpch", "orders"), "orderstatus"))), + new OutputColumnMetadata("totalprice", "double", ImmutableSet.of( + new SourceColumn(new QualifiedObjectName("iceberg", "tpch", "orders"), "totalprice"))), + new OutputColumnMetadata("orderdate", "date", ImmutableSet.of( + new SourceColumn(new QualifiedObjectName("iceberg", "tpch", "orders"), "orderdate"))), + new OutputColumnMetadata("orderpriority", "varchar", ImmutableSet.of( + new SourceColumn(new QualifiedObjectName("iceberg", "tpch", "orders"), "orderpriority"))), + new OutputColumnMetadata("clerk", "varchar", ImmutableSet.of( + new SourceColumn(new QualifiedObjectName("iceberg", "tpch", "orders"), "clerk"))), + new OutputColumnMetadata("shippriority", "integer", ImmutableSet.of( + new SourceColumn(new QualifiedObjectName("iceberg", "tpch", "orders"), "shippriority"))), + new OutputColumnMetadata("comment", "varchar", ImmutableSet.of( + new SourceColumn(new QualifiedObjectName("iceberg", "tpch", "orders"), "comment")))); + } + + @Test + public void testOutputColumnsForCreateTableAsSelectWithColumns() + throws Exception + { + runQueryAndWaitForEvents("CREATE TABLE create_update_table3 AS SELECT clerk, orderkey, totalprice FROM orders", 2); + QueryCompletedEvent event = getOnlyElement(generatedEvents.getQueryCompletedEvents()); + + assertThat(event.getIoMetadata().getOutput().get().getCatalogName()).isEqualTo("iceberg"); + assertThat(event.getIoMetadata().getOutput().get().getSchema()).isEqualTo("tpch"); + assertThat(event.getIoMetadata().getOutput().get().getTable()).isEqualTo("create_update_table3"); + assertThat(event.getMetadata().getUpdateQueryType().get()).isEqualTo("CREATE TABLE"); + + assertThat(event.getIoMetadata().getOutput().get().getColumns().get()) + .containsExactly( + new OutputColumnMetadata("clerk", "varchar", ImmutableSet.of( + new SourceColumn(new QualifiedObjectName("iceberg", "tpch", "orders"), "clerk"))), + new OutputColumnMetadata("orderkey", "bigint", ImmutableSet.of( + new SourceColumn(new QualifiedObjectName("iceberg", "tpch", "orders"), "orderkey"))), + new OutputColumnMetadata("totalprice", "double", ImmutableSet.of( + new SourceColumn(new QualifiedObjectName("iceberg", "tpch", "orders"), "totalprice")))); + } + + @Test + public void testOutputColumnsForCreateTableAsSelectWithAlias() + throws Exception + { + runQueryAndWaitForEvents("CREATE TABLE create_update_table4 AS SELECT clerk AS clerk_name, orderkey, totalprice AS annual_totalprice FROM orders", 2); + QueryCompletedEvent event = getOnlyElement(generatedEvents.getQueryCompletedEvents()); + + assertThat(event.getIoMetadata().getOutput().get().getCatalogName()).isEqualTo("iceberg"); + assertThat(event.getIoMetadata().getOutput().get().getSchema()).isEqualTo("tpch"); + assertThat(event.getIoMetadata().getOutput().get().getTable()).isEqualTo("create_update_table4"); + assertThat(event.getMetadata().getUpdateQueryType().get()).isEqualTo("CREATE TABLE"); + + assertThat(event.getIoMetadata().getOutput().get().getColumns().get()) + .containsExactly( + new OutputColumnMetadata("clerk_name", "varchar", ImmutableSet.of( + new SourceColumn(new QualifiedObjectName("iceberg", "tpch", "orders"), "clerk"))), + new OutputColumnMetadata("orderkey", "bigint", ImmutableSet.of( + new SourceColumn(new QualifiedObjectName("iceberg", "tpch", "orders"), "orderkey"))), + new OutputColumnMetadata("annual_totalprice", "double", ImmutableSet.of( + new SourceColumn(new QualifiedObjectName("iceberg", "tpch", "orders"), "totalprice")))); + } + + @Test + public void testOutputColumnsForCreateTableAsSelectAllWithAliasedRelation() + throws Exception + { + runQueryAndWaitForEvents("CREATE TABLE table_alias1 AS SELECT clerk AS test_clerk, orderkey AS test_orderkey, totalprice AS test_totalprice FROM (SELECT * from orders) orders_a", 2); + QueryCompletedEvent event = getOnlyElement(generatedEvents.getQueryCompletedEvents()); + + assertThat(event.getIoMetadata().getOutput().get().getCatalogName()).isEqualTo("iceberg"); + assertThat(event.getIoMetadata().getOutput().get().getSchema()).isEqualTo("tpch"); + assertThat(event.getIoMetadata().getOutput().get().getTable()).isEqualTo("table_alias1"); + assertThat(event.getMetadata().getUpdateQueryType().get()).isEqualTo("CREATE TABLE"); + + assertThat(event.getIoMetadata().getOutput().get().getColumns().get()) + .containsExactly( + new OutputColumnMetadata("test_clerk", "varchar", ImmutableSet.of( + new SourceColumn(new QualifiedObjectName("iceberg", "tpch", "orders"), "clerk"))), + new OutputColumnMetadata("test_orderkey", "bigint", ImmutableSet.of( + new SourceColumn(new QualifiedObjectName("iceberg", "tpch", "orders"), "orderkey"))), + new OutputColumnMetadata("test_totalprice", "double", ImmutableSet.of( + new SourceColumn(new QualifiedObjectName("iceberg", "tpch", "orders"), "totalprice")))); + } + + @Test + public void testOutputColumnsForCreateTableAsSelectColumnAliasInAliasedRelation() + throws Exception + { + runQueryAndWaitForEvents("CREATE TABLE table_alias2 AS SELECT aliased_clerk AS test_clerk, aliased_orderkey AS test_orderkey, aliased_totalprice AS test_totalprice FROM (SELECT clerk,orderkey,totalprice from orders) orders_a(aliased_clerk,aliased_orderkey,aliased_totalprice)", 2); QueryCompletedEvent event = getOnlyElement(generatedEvents.getQueryCompletedEvents()); assertThat(event.getIoMetadata().getOutput().get().getCatalogName()).isEqualTo("iceberg"); assertThat(event.getIoMetadata().getOutput().get().getSchema()).isEqualTo("tpch"); - assertThat(event.getIoMetadata().getOutput().get().getTable()).isEqualTo("create_update_table"); + assertThat(event.getIoMetadata().getOutput().get().getTable()).isEqualTo("table_alias2"); assertThat(event.getMetadata().getUpdateQueryType().get()).isEqualTo("CREATE TABLE"); assertThat(event.getIoMetadata().getOutput().get().getColumns().get()) .containsExactly( - new Column("orderkey", "bigint"), - new Column("custkey", "bigint"), - new Column("orderstatus", "varchar"), - new Column("totalprice", "double"), - new Column("orderdate", "date"), - new Column("orderpriority", "varchar"), - new Column("clerk", "varchar"), - new Column("shippriority", "integer"), - new Column("comment", "varchar")); + new OutputColumnMetadata("test_clerk", "varchar", ImmutableSet.of( + new SourceColumn(new QualifiedObjectName("iceberg", "tpch", "orders"), "clerk"))), + new OutputColumnMetadata("test_orderkey", "bigint", ImmutableSet.of( + new SourceColumn(new QualifiedObjectName("iceberg", "tpch", "orders"), "orderkey"))), + new OutputColumnMetadata("test_totalprice", "double", ImmutableSet.of( + new SourceColumn(new QualifiedObjectName("iceberg", "tpch", "orders"), "totalprice")))); + } + + @Test + public void testOutputColumnsForSetOperationUnion() + throws Exception + { + runQueryAndWaitForEvents("CREATE TABLE table_alias3 AS SELECT orderpriority AS test_orderpriority, orderkey AS test_orderkey FROM orders UNION SELECT clerk, custkey FROM orders", 2); + QueryCompletedEvent event = getOnlyElement(generatedEvents.getQueryCompletedEvents()); + + assertThat(event.getIoMetadata().getOutput().get().getColumns().get()) + .containsExactly( + new OutputColumnMetadata("test_orderpriority", "varchar", ImmutableSet.of( + new SourceColumn(new QualifiedObjectName("iceberg", "tpch", "orders"), "orderpriority"), + new SourceColumn(new QualifiedObjectName("iceberg", "tpch", "orders"), "clerk"))), + new OutputColumnMetadata("test_orderkey", "bigint", ImmutableSet.of( + new SourceColumn(new QualifiedObjectName("iceberg", "tpch", "orders"), "orderkey"), + new SourceColumn(new QualifiedObjectName("iceberg", "tpch", "orders"), "custkey")))); + } + + @Test + public void testOutputColumnsForSetOperationUnionAll() + throws Exception + { + runQueryAndWaitForEvents("CREATE TABLE table_alias4 AS SELECT orderpriority AS test_orderpriority, orderkey AS test_orderkey FROM orders UNION ALL SELECT clerk, custkey FROM orders", 2); + QueryCompletedEvent event = getOnlyElement(generatedEvents.getQueryCompletedEvents()); + + assertThat(event.getIoMetadata().getOutput().get().getColumns().get()) + .containsExactly( + new OutputColumnMetadata("test_orderpriority", "varchar", ImmutableSet.of( + new SourceColumn(new QualifiedObjectName("iceberg", "tpch", "orders"), "orderpriority"), + new SourceColumn(new QualifiedObjectName("iceberg", "tpch", "orders"), "clerk"))), + new OutputColumnMetadata("test_orderkey", "bigint", ImmutableSet.of( + new SourceColumn(new QualifiedObjectName("iceberg", "tpch", "orders"), "orderkey"), + new SourceColumn(new QualifiedObjectName("iceberg", "tpch", "orders"), "custkey")))); } static class TestingEventListenerPlugin diff --git a/presto-iceberg/src/test/java/com/facebook/presto/iceberg/hive/TestHiveTableOperationsConfig.java b/presto-iceberg/src/test/java/com/facebook/presto/iceberg/hive/TestHiveTableOperationsConfig.java index 5ac51d56b975b..42265f3fa5470 100644 --- a/presto-iceberg/src/test/java/com/facebook/presto/iceberg/hive/TestHiveTableOperationsConfig.java +++ b/presto-iceberg/src/test/java/com/facebook/presto/iceberg/hive/TestHiveTableOperationsConfig.java @@ -37,7 +37,8 @@ public void testDefaults() .setTableRefreshBackoffMaxSleepTime(succinctDuration(5, SECONDS)) .setTableRefreshMaxRetryTime(succinctDuration(1, MINUTES)) .setTableRefreshBackoffScaleFactor(4.0) - .setTableRefreshRetries(20)); + .setTableRefreshRetries(20) + .setLockingEnabled(true)); } @Test @@ -49,6 +50,7 @@ public void testExplicitPropertyMappings() .put("iceberg.hive.table-refresh.max-retry-time", "30s") .put("iceberg.hive.table-refresh.retries", "42") .put("iceberg.hive.table-refresh.backoff-scale-factor", "2.0") + .put("iceberg.engine.hive.lock-enabled", "false") .build(); IcebergHiveTableOperationsConfig expected = new IcebergHiveTableOperationsConfig() @@ -56,7 +58,8 @@ public void testExplicitPropertyMappings() .setTableRefreshBackoffMaxSleepTime(succinctDuration(20, SECONDS)) .setTableRefreshMaxRetryTime(succinctDuration(30, SECONDS)) .setTableRefreshBackoffScaleFactor(2.0) - .setTableRefreshRetries(42); + .setTableRefreshRetries(42) + .setLockingEnabled(false); assertFullMapping(properties, expected); } diff --git a/presto-iceberg/src/test/java/com/facebook/presto/iceberg/hive/TestIcebergDistributedHive.java b/presto-iceberg/src/test/java/com/facebook/presto/iceberg/hive/TestIcebergDistributedHive.java index f0ff6cf69b0c0..6e0697f4ebb65 100644 --- a/presto-iceberg/src/test/java/com/facebook/presto/iceberg/hive/TestIcebergDistributedHive.java +++ b/presto-iceberg/src/test/java/com/facebook/presto/iceberg/hive/TestIcebergDistributedHive.java @@ -16,7 +16,11 @@ import com.facebook.presto.Session; import com.facebook.presto.common.QualifiedObjectName; import com.facebook.presto.common.transaction.TransactionId; +import com.facebook.presto.hive.HdfsContext; +import com.facebook.presto.hive.HiveColumnConverterProvider; import com.facebook.presto.hive.metastore.ExtendedHiveMetastore; +import com.facebook.presto.hive.metastore.MetastoreContext; +import com.facebook.presto.iceberg.HiveTableOperations; import com.facebook.presto.iceberg.IcebergCatalogName; import com.facebook.presto.iceberg.IcebergDistributedTestBase; import com.facebook.presto.iceberg.IcebergHiveMetadata; @@ -26,7 +30,9 @@ import com.facebook.presto.metadata.CatalogManager; import com.facebook.presto.metadata.CatalogMetadata; import com.facebook.presto.metadata.MetadataUtil; +import com.facebook.presto.spi.ColumnMetadata; import com.facebook.presto.spi.ConnectorId; +import com.facebook.presto.spi.ConnectorSession; import com.facebook.presto.spi.SchemaTableName; import com.facebook.presto.spi.TableHandle; import com.facebook.presto.spi.connector.classloader.ClassLoaderSafeConnectorMetadata; @@ -34,20 +40,33 @@ import com.google.common.cache.CacheBuilder; import com.google.common.cache.CacheStats; import com.google.common.collect.ImmutableMap; +import org.apache.iceberg.BaseTable; +import org.apache.iceberg.PartitionSpec; import org.apache.iceberg.Table; +import org.apache.iceberg.TableMetadata; +import org.apache.iceberg.Transaction; import org.testng.annotations.Test; import java.lang.reflect.Field; +import java.util.ArrayList; import java.util.HashMap; +import java.util.List; import java.util.Map; import java.util.Optional; +import static com.facebook.presto.common.type.IntegerType.INTEGER; import static com.facebook.presto.hive.metastore.InMemoryCachingHiveMetastore.memoizeMetastore; +import static com.facebook.presto.hive.metastore.MetastoreUtil.getMetastoreHeaders; +import static com.facebook.presto.hive.metastore.MetastoreUtil.isUserDefinedTypeEncodingEnabled; import static com.facebook.presto.iceberg.CatalogType.HIVE; +import static com.facebook.presto.iceberg.IcebergAbstractMetadata.toIcebergSchema; import static com.facebook.presto.iceberg.IcebergQueryRunner.ICEBERG_CATALOG; import static com.facebook.presto.spi.statistics.ColumnStatisticType.NUMBER_OF_DISTINCT_VALUES; import static com.facebook.presto.spi.statistics.ColumnStatisticType.TOTAL_SIZE_IN_BYTES; +import static com.google.common.io.Files.createTempDir; import static java.lang.String.format; +import static org.apache.iceberg.TableMetadata.newTableMetadata; +import static org.apache.iceberg.Transactions.createTableTransaction; import static org.testng.Assert.assertEquals; import static org.testng.Assert.assertFalse; import static org.testng.Assert.assertTrue; @@ -56,6 +75,15 @@ public class TestIcebergDistributedHive extends IcebergDistributedTestBase { + public TestIcebergDistributedHive(Map extraConnectorProperties) + { + super(HIVE, ImmutableMap.builder() + .put("iceberg.hive-statistics-merge-strategy", Joiner.on(",").join( + NUMBER_OF_DISTINCT_VALUES.name(), + TOTAL_SIZE_IN_BYTES.name())) + .putAll(extraConnectorProperties) + .build()); + } public TestIcebergDistributedHive() { super(HIVE, ImmutableMap.of("iceberg.hive-statistics-merge-strategy", Joiner.on(",").join(NUMBER_OF_DISTINCT_VALUES.name(), TOTAL_SIZE_IN_BYTES.name()))); @@ -187,6 +215,21 @@ public void testManifestFileCachingDisabled() assertQuerySucceeds(session, "DROP SCHEMA default"); } + @Test + public void testCommitTableMetadataForNoLock() + { + createTable("iceberg-test-table", createTempDir().toURI().toString(), ImmutableMap.of("engine.hive.lock-enabled", "false"), 2); + BaseTable table = (BaseTable) loadTable("iceberg-test-table"); + HiveTableOperations operations = (HiveTableOperations) table.operations(); + TableMetadata currentMetadata = operations.current(); + + ImmutableMap.Builder builder = ImmutableMap.builder(); + builder.putAll(currentMetadata.properties()); + builder.put("test_property_new", "test_value_new"); + operations.commit(currentMetadata, TableMetadata.buildFrom(currentMetadata).setProperties(builder.build()).build()); + assertEquals(operations.current().properties(), builder.build()); + } + @Override protected Table loadTable(String tableName) { @@ -210,4 +253,35 @@ protected ExtendedHiveMetastore getFileHiveMetastore() "test"); return memoizeMetastore(fileHiveMetastore, false, 1000, 0); } + + protected Table createTable(String tableName, String targetPath, Map tableProperties, int columns) + { + CatalogManager catalogManager = getDistributedQueryRunner().getCoordinator().getCatalogManager(); + ConnectorId connectorId = catalogManager.getCatalog(getDistributedQueryRunner().getDefaultSession().getCatalog().get()).get().getConnectorId(); + ConnectorSession session = getQueryRunner().getDefaultSession().toConnectorSession(connectorId); + MetastoreContext context = new MetastoreContext(session.getIdentity(), session.getQueryId(), session.getClientInfo(), session.getClientTags(), session.getSource(), getMetastoreHeaders(session), isUserDefinedTypeEncodingEnabled(session), HiveColumnConverterProvider.DEFAULT_COLUMN_CONVERTER_PROVIDER, session.getWarningCollector(), session.getRuntimeStats()); + HdfsContext hdfsContext = new HdfsContext(session, "tpch", tableName); + HiveTableOperations operations = new HiveTableOperations( + getFileHiveMetastore(), + context, + getHdfsEnvironment(), + hdfsContext, + new IcebergHiveTableOperationsConfig().setLockingEnabled(false), + new ManifestFileCache(CacheBuilder.newBuilder().build(), false, 0, 1024), + "tpch", + tableName, + session.getUser(), + targetPath); + List columnMetadataList = new ArrayList<>(); + for (int i = 0; i < columns; i++) { + columnMetadataList.add(ColumnMetadata.builder().setName("column" + i).setType(INTEGER).build()); + } + TableMetadata metadata = newTableMetadata( + toIcebergSchema(columnMetadataList), + PartitionSpec.unpartitioned(), targetPath, + tableProperties); + Transaction transaction = createTableTransaction(tableName, operations, metadata); + transaction.commitTransaction(); + return transaction.table(); + } } diff --git a/presto-iceberg/src/test/java/org/apache/iceberg/rest/IcebergRestCatalogServlet.java b/presto-iceberg/src/test/java/org/apache/iceberg/rest/IcebergRestCatalogServlet.java index 588f556065616..0fca1ffae07aa 100644 --- a/presto-iceberg/src/test/java/org/apache/iceberg/rest/IcebergRestCatalogServlet.java +++ b/presto-iceberg/src/test/java/org/apache/iceberg/rest/IcebergRestCatalogServlet.java @@ -25,7 +25,7 @@ import org.apache.iceberg.exceptions.RESTException; import org.apache.iceberg.relocated.com.google.common.collect.ImmutableMap; import org.apache.iceberg.relocated.com.google.common.io.CharStreams; -import org.apache.iceberg.rest.RESTCatalogAdapter.HTTPMethod; +import org.apache.iceberg.rest.HTTPRequest.HTTPMethod; import org.apache.iceberg.rest.RESTCatalogAdapter.Route; import org.apache.iceberg.rest.responses.ErrorResponse; import org.apache.iceberg.util.Pair; @@ -111,15 +111,18 @@ protected void execute(ServletRequestContext context, HttpServletResponse respon } try { + HTTPRequest request = restCatalogAdapter.buildRequest( + context.method(), + context.path(), + context.queryParams(), + context.headers(), + context.body()); Object responseBody = restCatalogAdapter.execute( - context.method(), - context.path(), - context.queryParams(), - context.body(), + request, context.route().responseClass(), - context.headers(), - handle(response)); + handleResponseError(response), + handleResponseHeader(response)); if (responseBody != null) { RESTObjectMapper.mapper().writeValue(response.getWriter(), responseBody); @@ -143,7 +146,15 @@ protected void execute(ServletRequestContext context, HttpServletResponse respon } } - protected Consumer handle(HttpServletResponse response) + private Consumer> handleResponseHeader(HttpServletResponse response) + { + return (responseHeaders) -> { + LOG.error("Unexpected response header: %s", responseHeaders); + throw new RuntimeException("Unexpected response header: " + responseHeaders); + }; + } + + protected Consumer handleResponseError(HttpServletResponse response) { return (errorResponse) -> { response.setStatus(errorResponse.code()); diff --git a/presto-jdbc/pom.xml b/presto-jdbc/pom.xml index c4fc6237d602c..d5515e0d71b7c 100644 --- a/presto-jdbc/pom.xml +++ b/presto-jdbc/pom.xml @@ -117,6 +117,30 @@ + + com.facebook.airlift + concurrent + test + + + + com.facebook.airlift + configuration + test + + + + com.facebook.airlift + http-server + test + + + + com.facebook.airlift + jaxrs + test + + com.facebook.presto presto-testng-services @@ -177,6 +201,18 @@ test + + jakarta.servlet + jakarta.servlet-api + test + + + + jakarta.ws.rs + jakarta.ws.rs-api + test + + org.testng testng @@ -201,12 +237,6 @@ test - - com.facebook.airlift - concurrent - test - - com.google.inject guice diff --git a/presto-jdbc/src/main/java/com/facebook/presto/jdbc/ConnectionProperties.java b/presto-jdbc/src/main/java/com/facebook/presto/jdbc/ConnectionProperties.java index c6c8e180820d7..931761723d9e5 100644 --- a/presto-jdbc/src/main/java/com/facebook/presto/jdbc/ConnectionProperties.java +++ b/presto-jdbc/src/main/java/com/facebook/presto/jdbc/ConnectionProperties.java @@ -13,6 +13,9 @@ */ package com.facebook.presto.jdbc; +import com.facebook.airlift.units.Duration; +import com.facebook.presto.client.auth.external.ExternalRedirectStrategy; +import com.google.common.base.Splitter; import com.google.common.collect.ImmutableMap; import com.google.common.collect.ImmutableSet; import com.google.common.net.HostAndPort; @@ -26,12 +29,14 @@ import java.util.Properties; import java.util.Set; import java.util.function.Predicate; +import java.util.stream.StreamSupport; import static com.facebook.presto.jdbc.AbstractConnectionProperty.ClassListConverter.CLASS_LIST_CONVERTER; import static com.facebook.presto.jdbc.AbstractConnectionProperty.HttpProtocolConverter.HTTP_PROTOCOL_CONVERTER; import static com.facebook.presto.jdbc.AbstractConnectionProperty.ListValidateConvertor.LIST_VALIDATE_CONVERTOR; import static com.facebook.presto.jdbc.AbstractConnectionProperty.StringMapConverter.STRING_MAP_CONVERTER; import static com.facebook.presto.jdbc.AbstractConnectionProperty.checkedPredicate; +import static com.google.common.collect.ImmutableList.toImmutableList; import static java.util.Collections.unmodifiableMap; import static java.util.function.Function.identity; import static java.util.stream.Collectors.toMap; @@ -68,6 +73,11 @@ final class ConnectionProperties public static final ConnectionProperty FOLLOW_REDIRECTS = new FollowRedirects(); public static final ConnectionProperty SSL_KEY_STORE_TYPE = new SSLKeyStoreType(); public static final ConnectionProperty SSL_TRUST_STORE_TYPE = new SSLTrustStoreType(); + public static final ConnectionProperty EXTERNAL_AUTHENTICATION = new ExternalAuthentication(); + public static final ConnectionProperty EXTERNAL_AUTHENTICATION_TIMEOUT = new ExternalAuthenticationTimeout(); + public static final ConnectionProperty EXTERNAL_AUTHENTICATION_TOKEN_CACHE = new ExternalAuthenticationTokenCache(); + public static final ConnectionProperty> EXTERNAL_AUTHENTICATION_REDIRECT_HANDLERS = new ExternalAuthenticationRedirectHandlers(); + private static final Set> ALL_PROPERTIES = ImmutableSet.>builder() .add(USER) .add(PASSWORD) @@ -98,6 +108,10 @@ final class ConnectionProperties .add(QUERY_INTERCEPTORS) .add(VALIDATE_NEXTURI_SOURCE) .add(FOLLOW_REDIRECTS) + .add(EXTERNAL_AUTHENTICATION) + .add(EXTERNAL_AUTHENTICATION_TIMEOUT) + .add(EXTERNAL_AUTHENTICATION_TOKEN_CACHE) + .add(EXTERNAL_AUTHENTICATION_REDIRECT_HANDLERS) .build(); private static final Map> KEY_LOOKUP = unmodifiableMap(ALL_PROPERTIES.stream() @@ -411,4 +425,54 @@ public SSLKeyStoreType() super("SSLKeyStoreType", Optional.of(KeyStore.getDefaultType()), NOT_REQUIRED, ALLOWED, STRING_CONVERTER); } } + + private static Predicate isExternalAuthEnabled() + { + return checkedPredicate(properties -> EXTERNAL_AUTHENTICATION.getValue(properties).isPresent()); + } + + private static class ExternalAuthentication + extends AbstractConnectionProperty + { + public ExternalAuthentication() + { + super("externalAuthentication", Optional.of("false"), NOT_REQUIRED, ALLOWED, BOOLEAN_CONVERTER); + } + } + + private static class ExternalAuthenticationTimeout + extends AbstractConnectionProperty + { + public ExternalAuthenticationTimeout() + { + super("externalAuthenticationTimeout", NOT_REQUIRED, isExternalAuthEnabled(), Duration::valueOf); + } + } + + private static class ExternalAuthenticationTokenCache + extends AbstractConnectionProperty + { + public ExternalAuthenticationTokenCache() + { + super("externalAuthenticationTokenCache", Optional.of(KnownTokenCache.NONE.name()), NOT_REQUIRED, ALLOWED, KnownTokenCache::valueOf); + } + } + + private static class ExternalAuthenticationRedirectHandlers + extends AbstractConnectionProperty> + { + private static final Splitter ENUM_SPLITTER = Splitter.on(',').trimResults().omitEmptyStrings(); + + public ExternalAuthenticationRedirectHandlers() + { + super("externalAuthenticationRedirectHandlers", Optional.of("OPEN"), NOT_REQUIRED, ALLOWED, ExternalAuthenticationRedirectHandlers::parse); + } + + public static List parse(String value) + { + return StreamSupport.stream(ENUM_SPLITTER.split(value).spliterator(), false) + .map(ExternalRedirectStrategy::valueOf) + .collect(toImmutableList()); + } + } } diff --git a/presto-jdbc/src/main/java/com/facebook/presto/jdbc/KnownTokenCache.java b/presto-jdbc/src/main/java/com/facebook/presto/jdbc/KnownTokenCache.java new file mode 100644 index 0000000000000..8792075be56d1 --- /dev/null +++ b/presto-jdbc/src/main/java/com/facebook/presto/jdbc/KnownTokenCache.java @@ -0,0 +1,29 @@ +/* + * Licensed 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 com.facebook.presto.jdbc; + +import com.facebook.presto.client.auth.external.KnownToken; + +public enum KnownTokenCache +{ + NONE { + @Override + KnownToken create() + { + return KnownToken.local(); + } + }; + + abstract KnownToken create(); +} diff --git a/presto-jdbc/src/main/java/com/facebook/presto/jdbc/PrestoDriverUri.java b/presto-jdbc/src/main/java/com/facebook/presto/jdbc/PrestoDriverUri.java index e355b58f32331..1dcfc47b152d8 100644 --- a/presto-jdbc/src/main/java/com/facebook/presto/jdbc/PrestoDriverUri.java +++ b/presto-jdbc/src/main/java/com/facebook/presto/jdbc/PrestoDriverUri.java @@ -15,6 +15,12 @@ import com.facebook.presto.client.ClientException; import com.facebook.presto.client.OkHttpUtil; +import com.facebook.presto.client.auth.external.CompositeRedirectHandler; +import com.facebook.presto.client.auth.external.ExternalAuthenticator; +import com.facebook.presto.client.auth.external.HttpTokenPoller; +import com.facebook.presto.client.auth.external.RedirectHandler; +import com.facebook.presto.client.auth.external.TokenPoller; +import com.google.common.annotations.VisibleForTesting; import com.google.common.base.Splitter; import com.google.common.collect.ImmutableList; import com.google.common.collect.ImmutableMap; @@ -27,6 +33,7 @@ import java.net.URI; import java.net.URISyntaxException; import java.sql.SQLException; +import java.time.Duration; import java.util.Arrays; import java.util.HashMap; import java.util.List; @@ -35,6 +42,7 @@ import java.util.Optional; import java.util.Properties; import java.util.TimeZone; +import java.util.concurrent.atomic.AtomicReference; import static com.facebook.presto.client.GCSOAuthInterceptor.GCS_CREDENTIALS_PATH_KEY; import static com.facebook.presto.client.GCSOAuthInterceptor.GCS_OAUTH_SCOPES_KEY; @@ -51,6 +59,10 @@ import static com.facebook.presto.jdbc.ConnectionProperties.CLIENT_TAGS; import static com.facebook.presto.jdbc.ConnectionProperties.CUSTOM_HEADERS; import static com.facebook.presto.jdbc.ConnectionProperties.DISABLE_COMPRESSION; +import static com.facebook.presto.jdbc.ConnectionProperties.EXTERNAL_AUTHENTICATION; +import static com.facebook.presto.jdbc.ConnectionProperties.EXTERNAL_AUTHENTICATION_REDIRECT_HANDLERS; +import static com.facebook.presto.jdbc.ConnectionProperties.EXTERNAL_AUTHENTICATION_TIMEOUT; +import static com.facebook.presto.jdbc.ConnectionProperties.EXTERNAL_AUTHENTICATION_TOKEN_CACHE; import static com.facebook.presto.jdbc.ConnectionProperties.EXTRA_CREDENTIALS; import static com.facebook.presto.jdbc.ConnectionProperties.FOLLOW_REDIRECTS; import static com.facebook.presto.jdbc.ConnectionProperties.HTTP_PROTOCOLS; @@ -88,7 +100,7 @@ final class PrestoDriverUri private static final Splitter QUERY_SPLITTER = Splitter.on('&').omitEmptyStrings(); private static final Splitter ARG_SPLITTER = Splitter.on('=').limit(2); - + private static final AtomicReference REDIRECT_HANDLER = new AtomicReference<>(null); private final HostAndPort address; private final URI uri; @@ -282,6 +294,30 @@ public void setupClient(OkHttpClient.Builder builder) } builder.addInterceptor(tokenAuth(ACCESS_TOKEN.getValue(properties).get())); } + + if (EXTERNAL_AUTHENTICATION.getValue(properties).orElse(false)) { + if (!useSecureConnection) { + throw new SQLException("Authentication using external authorization requires SSL to be enabled"); + } + + // create HTTP client that shares the same settings, but without the external authenticator + TokenPoller poller = new HttpTokenPoller(builder.build()); + + Duration timeout = EXTERNAL_AUTHENTICATION_TIMEOUT.getValue(properties) + .map(value -> Duration.ofMillis(value.toMillis())) + .orElse(Duration.ofMinutes(2)); + + KnownTokenCache knownTokenCache = EXTERNAL_AUTHENTICATION_TOKEN_CACHE.getValue(properties).get(); + Optional configuredHandler = EXTERNAL_AUTHENTICATION_REDIRECT_HANDLERS.getValue(properties) + .map(CompositeRedirectHandler::new) + .map(RedirectHandler.class::cast); + RedirectHandler redirectHandler = Optional.ofNullable(REDIRECT_HANDLER.get()) + .orElseGet(() -> configuredHandler.orElseThrow(() -> new RuntimeException("External authentication redirect handler is not configured"))); + ExternalAuthenticator authenticator = new ExternalAuthenticator(redirectHandler, poller, knownTokenCache.create(), timeout); + + builder.authenticator(authenticator); + builder.addInterceptor(authenticator); + } } catch (ClientException e) { throw new SQLException(e.getMessage(), e); @@ -419,4 +455,10 @@ private static void validateConnectionProperties(Properties connectionProperties property.validate(connectionProperties); } } + + @VisibleForTesting + static void setRedirectHandler(RedirectHandler handler) + { + REDIRECT_HANDLER.set(requireNonNull(handler, "handler is null")); + } } diff --git a/presto-jdbc/src/test/java/com/facebook/presto/jdbc/TestJdbcExternalAuthentication.java b/presto-jdbc/src/test/java/com/facebook/presto/jdbc/TestJdbcExternalAuthentication.java new file mode 100644 index 0000000000000..ede9a534cf2a3 --- /dev/null +++ b/presto-jdbc/src/test/java/com/facebook/presto/jdbc/TestJdbcExternalAuthentication.java @@ -0,0 +1,524 @@ +/* + * Licensed 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 com.facebook.presto.jdbc; + +import com.facebook.airlift.configuration.AbstractConfigurationAwareModule; +import com.facebook.airlift.http.server.AuthenticationException; +import com.facebook.airlift.http.server.Authenticator; +import com.facebook.airlift.http.server.BasicPrincipal; +import com.facebook.airlift.log.Logging; +import com.facebook.presto.client.ClientException; +import com.facebook.presto.client.auth.external.DesktopBrowserRedirectHandler; +import com.facebook.presto.client.auth.external.RedirectException; +import com.facebook.presto.client.auth.external.RedirectHandler; +import com.facebook.presto.server.security.SecurityConfig; +import com.facebook.presto.server.testing.TestingPrestoServer; +import com.facebook.presto.sql.parser.SqlParserOptions; +import com.facebook.presto.tpch.TpchPlugin; +import com.google.common.collect.ImmutableList; +import com.google.common.collect.ImmutableMap; +import com.google.inject.Binder; +import com.google.inject.Inject; +import com.google.inject.Key; +import com.google.inject.Module; +import com.google.inject.Scopes; +import com.google.inject.multibindings.Multibinder; +import jakarta.servlet.http.HttpServletRequest; +import jakarta.ws.rs.GET; +import jakarta.ws.rs.Path; +import jakarta.ws.rs.PathParam; +import jakarta.ws.rs.Produces; +import jakarta.ws.rs.core.Context; +import jakarta.ws.rs.core.Response; +import okhttp3.HttpUrl; +import okhttp3.OkHttpClient; +import okhttp3.Request; +import org.testng.annotations.AfterClass; +import org.testng.annotations.BeforeClass; +import org.testng.annotations.BeforeMethod; +import org.testng.annotations.Test; + +import java.io.File; +import java.io.IOException; +import java.net.URI; +import java.security.Principal; +import java.sql.Connection; +import java.sql.DriverManager; +import java.sql.SQLException; +import java.sql.Statement; +import java.util.ConcurrentModificationException; +import java.util.List; +import java.util.Map; +import java.util.Optional; +import java.util.Properties; +import java.util.Set; +import java.util.UUID; +import java.util.concurrent.ConcurrentHashMap; +import java.util.concurrent.atomic.AtomicReference; +import java.util.function.IntSupplier; + +import static com.facebook.airlift.configuration.ConditionalModule.installModuleIf; +import static com.facebook.airlift.jaxrs.JaxrsBinder.jaxrsBinder; +import static com.facebook.airlift.testing.Closeables.closeAll; +import static com.facebook.presto.jdbc.PrestoDriverUri.setRedirectHandler; +import static com.facebook.presto.jdbc.TestPrestoDriver.waitForNodeRefresh; +import static com.google.common.io.Resources.getResource; +import static com.google.inject.Scopes.SINGLETON; +import static com.google.inject.multibindings.Multibinder.newSetBinder; +import static com.google.inject.util.Modules.combine; +import static jakarta.ws.rs.core.HttpHeaders.AUTHORIZATION; +import static jakarta.ws.rs.core.MediaType.APPLICATION_JSON_TYPE; +import static jakarta.ws.rs.core.MediaType.TEXT_PLAIN; +import static jakarta.ws.rs.core.Response.Status.NOT_FOUND; +import static java.lang.String.format; +import static java.net.HttpURLConnection.HTTP_OK; +import static java.util.Objects.requireNonNull; +import static org.assertj.core.api.Assertions.assertThat; +import static org.assertj.core.api.Assertions.assertThatThrownBy; + +@Test(singleThreaded = true) +public class TestJdbcExternalAuthentication +{ + private static final String TEST_CATALOG = "test_catalog"; + private TestingPrestoServer server; + + @BeforeClass + public void setup() + throws Exception + { + Logging.initialize(); + + Map properties = ImmutableMap.builder() + .put("http-server.authentication.type", "TEST_EXTERNAL") + .put("http-server.https.enabled", "true") + .put("http-server.https.keystore.path", new File(getResource("localhost.keystore").toURI()).getPath()) + .put("http-server.https.keystore.key", "changeit") + .build(); + List additionalModules = ImmutableList.builder() + .add(new DummyExternalAuthModule(() -> server.getAddress().getPort())) + .build(); + + server = new TestingPrestoServer(true, properties, null, null, new SqlParserOptions(), additionalModules); + server.installPlugin(new TpchPlugin()); + server.createCatalog(TEST_CATALOG, "tpch"); + waitForNodeRefresh(server); + } + + @AfterClass(alwaysRun = true) + public void teardown() + throws Exception + { + closeAll(server); + server = null; + } + + @BeforeMethod(alwaysRun = true) + public void clearUpLoggingSessions() + { + invalidateAllTokens(); + } + + @Test + public void testSuccessfulAuthenticationWithHttpGetOnlyRedirectHandler() + throws Exception + { + try (RedirectHandlerFixture ignore = RedirectHandlerFixture.withHandler(new HttpGetOnlyRedirectHandler()); + Connection connection = createConnection(); + Statement statement = connection.createStatement()) { + assertThat(statement.execute("SELECT 123")).isTrue(); + } + } + + /** + * Ignored due to lack of ui environment with web-browser on CI servers. + * Still this test is useful for local environments. + */ + @Test(enabled = false) + public void testSuccessfulAuthenticationWithDefaultBrowserRedirect() + throws Exception + { + try (Connection connection = createConnection(); + Statement statement = connection.createStatement()) { + assertThat(statement.execute("SELECT 123")).isTrue(); + } + } + + @Test + public void testAuthenticationFailsAfterUnfinishedRedirect() + throws Exception + { + try (RedirectHandlerFixture ignore = RedirectHandlerFixture.withHandler(new NoOpRedirectHandler()); + Connection connection = createConnection(); + Statement statement = connection.createStatement()) { + assertThatThrownBy(() -> statement.execute("SELECT 123")) + .isInstanceOf(SQLException.class); + } + } + + @Test + public void testAuthenticationFailsAfterRedirectException() + throws Exception + { + try (RedirectHandlerFixture ignore = RedirectHandlerFixture.withHandler(new FailingRedirectHandler()); + Connection connection = createConnection(); + Statement statement = connection.createStatement()) { + assertThatThrownBy(() -> statement.execute("SELECT 123")) + .isInstanceOf(SQLException.class) + .hasCauseExactlyInstanceOf(RedirectException.class); + } + } + + @Test + public void testAuthenticationFailsAfterServerAuthenticationFailure() + throws Exception + { + try (RedirectHandlerFixture ignore = RedirectHandlerFixture.withHandler(new HttpGetOnlyRedirectHandler()); + AutoCloseable ignore2 = TokenPollingErrorFixture.withPollingError("error occurred during token polling"); + Connection connection = createConnection(); + Statement statement = connection.createStatement()) { + assertThatThrownBy(() -> statement.execute("SELECT 123")) + .isInstanceOf(SQLException.class) + .hasMessage("error occurred during token polling"); + } + } + + @Test + public void testAuthenticationFailsAfterReceivingMalformedHeaderFromServer() + throws Exception + { + try (RedirectHandlerFixture ignore = RedirectHandlerFixture.withHandler(new HttpGetOnlyRedirectHandler()); + AutoCloseable ignored = WwwAuthenticateHeaderFixture.withWwwAuthenticate("Bearer no-valid-fields"); + Connection connection = createConnection(); + Statement statement = connection.createStatement()) { + assertThatThrownBy(() -> statement.execute("SELECT 123")) + .isInstanceOf(SQLException.class) + .hasCauseInstanceOf(ClientException.class) + .hasMessage("Authentication failed: Authentication required"); + } + } + + @Test + public void testAuthenticationReusesObtainedTokenPerConnection() + throws Exception + { + try (RedirectHandlerFixture ignore = RedirectHandlerFixture.withHandler(new HttpGetOnlyRedirectHandler()); + Connection connection = createConnection(); + Statement statement = connection.createStatement()) { + statement.execute("SELECT 123"); + statement.execute("SELECT 123"); + statement.execute("SELECT 123"); + + assertThat(countIssuedTokens()).isEqualTo(1); + } + } + + @Test + public void testAuthenticationAfterInitialTokenHasBeenInvalidated() + throws Exception + { + try (RedirectHandlerFixture ignore = RedirectHandlerFixture.withHandler(new HttpGetOnlyRedirectHandler()); + Connection connection = createConnection(); + Statement statement = connection.createStatement()) { + statement.execute("SELECT 123"); + + invalidateAllTokens(); + assertThat(countIssuedTokens()).isEqualTo(0); + + assertThat(statement.execute("SELECT 123")).isTrue(); + } + } + + private Connection createConnection() + throws Exception + { + String url = format("jdbc:presto://localhost:%s", server.getHttpsAddress().getPort()); + Properties properties = new Properties(); + properties.setProperty("SSL", "true"); + properties.setProperty("SSLTrustStorePath", new File(getResource("localhost.truststore").toURI()).getPath()); + properties.setProperty("SSLTrustStorePassword", "changeit"); + properties.setProperty("externalAuthentication", "true"); + properties.setProperty("externalAuthenticationTimeout", "2s"); + properties.setProperty("user", "test"); + return DriverManager.getConnection(url, properties); + } + + private static Multibinder authenticatorBinder(Binder binder) + { + return newSetBinder(binder, Authenticator.class); + } + + public static Module authenticatorModule(Class clazz, Module module) + { + Module authModule = binder -> authenticatorBinder(binder).addBinding().to(clazz).in(Scopes.SINGLETON); + return installModuleIf( + SecurityConfig.class, + config -> true, + combine(module, authModule)); + } + + private static class DummyExternalAuthModule + extends AbstractConfigurationAwareModule + { + private final IntSupplier port; + + public DummyExternalAuthModule(IntSupplier port) + { + this.port = requireNonNull(port, "port is null"); + } + + @Override + protected void setup(Binder ignored) + { + Module test = authenticatorModule(DummyAuthenticator.class, binder -> { + binder.bind(Authentications.class).in(SINGLETON); + binder.bind(IntSupplier.class).toInstance(port); + jaxrsBinder(binder).bind(DummyExternalAuthResources.class); + }); + + install(test); + } + } + + private static class Authentications + { + private final Map logginSessions = new ConcurrentHashMap<>(); + private final Set validTokens = ConcurrentHashMap.newKeySet(); + + public String startAuthentication() + { + String sessionId = UUID.randomUUID().toString(); + logginSessions.put(sessionId, ""); + return sessionId; + } + + public void logIn(String sessionId) + { + String token = sessionId + "_token"; + validTokens.add(token); + logginSessions.put(sessionId, token); + } + + public Optional getToken(String sessionId) + throws IllegalArgumentException + { + return Optional.ofNullable(logginSessions.get(sessionId)) + .filter(s -> !s.isEmpty()); + } + + public boolean verifyToken(String token) + { + return validTokens.contains(token); + } + + public void invalidateAllTokens() + { + validTokens.clear(); + } + + public int countValidTokens() + { + return validTokens.size(); + } + } + + private void invalidateAllTokens() + { + Authentications authentications = server.getInstance(Key.get(Authentications.class)); + authentications.invalidateAllTokens(); + } + + private int countIssuedTokens() + { + Authentications authentications = server.getInstance(Key.get(Authentications.class)); + return authentications.countValidTokens(); + } + + public static class DummyAuthenticator + implements Authenticator + { + private final IntSupplier port; + private final Authentications authentications; + + @Inject + public DummyAuthenticator(IntSupplier port, Authentications authentications) + { + this.port = requireNonNull(port, "port is null"); + this.authentications = requireNonNull(authentications, "authentications is null"); + } + + @Override + public Principal authenticate(HttpServletRequest request) + throws AuthenticationException + { + Optional authHeader = Optional.ofNullable(request.getHeader(AUTHORIZATION)); + List bearerHeaders = authHeader.isPresent() ? ImmutableList.of(authHeader.get()) : ImmutableList.of(); + if (bearerHeaders.stream() + .filter(header -> header.startsWith("Bearer ")) + .anyMatch(header -> authentications.verifyToken(header.substring("Bearer ".length())))) { + return new BasicPrincipal("user"); + } + + String sessionId = authentications.startAuthentication(); + + throw Optional.ofNullable(WwwAuthenticateHeaderFixture.HEADER.get()) + .map(header -> new AuthenticationException("Authentication required", header)) + .orElseGet(() -> new AuthenticationException( + "Authentication required", + format("Bearer x_redirect_server=\"http://localhost:%s/v1/authentications/dummy/logins/%s\", " + + "x_token_server=\"http://localhost:%s/v1/authentications/dummy/%s\"", + port.getAsInt(), sessionId, port.getAsInt(), sessionId))); + } + } + + @Path("/v1/authentications/dummy") + public static class DummyExternalAuthResources + { + private final Authentications authentications; + + @Inject + public DummyExternalAuthResources(Authentications authentications) + { + this.authentications = authentications; + } + + @GET + @Produces(TEXT_PLAIN) + @Path("logins/{sessionId}") + public String logInUser(@PathParam("sessionId") String sessionId) + { + authentications.logIn(sessionId); + return "User has been successfully logged in during " + sessionId + " session"; + } + + @GET + @Path("{sessionId}") + public Response getToken(@PathParam("sessionId") String sessionId, @Context HttpServletRequest request) + { + try { + return Optional.ofNullable(TokenPollingErrorFixture.ERROR.get()) + .map(error -> Response.ok(format("{ \"error\" : \"%s\"}", error), APPLICATION_JSON_TYPE).build()) + .orElseGet(() -> authentications.getToken(sessionId) + .map(token -> Response.ok(format("{ \"token\" : \"%s\"}", token), APPLICATION_JSON_TYPE).build()) + .orElseGet(() -> Response.ok(format("{ \"nextUri\" : \"%s\" }", request.getRequestURI()), APPLICATION_JSON_TYPE).build())); + } + catch (IllegalArgumentException ex) { + return Response.status(NOT_FOUND).build(); + } + } + } + + public static class HttpGetOnlyRedirectHandler + implements RedirectHandler + { + @Override + public void redirectTo(URI uri) + throws RedirectException + { + OkHttpClient client = new OkHttpClient(); + + Request request = new Request.Builder() + .url(HttpUrl.get(uri.toString())) + .build(); + + try (okhttp3.Response response = client.newCall(request).execute()) { + if (response.code() != HTTP_OK) { + throw new RedirectException("HTTP GET failed with status " + response.code()); + } + } + catch (IOException e) { + throw new RedirectException("Redirection failed", e); + } + } + } + + public static class NoOpRedirectHandler + implements RedirectHandler + { + @Override + public void redirectTo(URI uri) + throws RedirectException + {} + } + + public static class FailingRedirectHandler + implements RedirectHandler + { + @Override + public void redirectTo(URI uri) + throws RedirectException + { + throw new RedirectException("Redirect to uri has failed " + uri); + } + } + + static class RedirectHandlerFixture + implements AutoCloseable + { + private static final RedirectHandlerFixture INSTANCE = new RedirectHandlerFixture(); + + private RedirectHandlerFixture() {} + + public static RedirectHandlerFixture withHandler(RedirectHandler handler) + { + setRedirectHandler(handler); + return INSTANCE; + } + + @Override + public void close() + { + setRedirectHandler(new DesktopBrowserRedirectHandler()); + } + } + + static class TokenPollingErrorFixture + implements AutoCloseable + { + private static final AtomicReference ERROR = new AtomicReference<>(null); + + public static AutoCloseable withPollingError(String error) + { + if (ERROR.compareAndSet(null, error)) { + return new TokenPollingErrorFixture(); + } + throw new ConcurrentModificationException("polling errors can't be invoked in parallel"); + } + + @Override + public void close() + { + ERROR.set(null); + } + } + + static class WwwAuthenticateHeaderFixture + implements AutoCloseable + { + private static final AtomicReference HEADER = new AtomicReference<>(null); + + public static AutoCloseable withWwwAuthenticate(String header) + { + if (HEADER.compareAndSet(null, header)) { + return new WwwAuthenticateHeaderFixture(); + } + throw new ConcurrentModificationException("with WWW-Authenticate header can't be invoked in parallel"); + } + + @Override + public void close() + { + HEADER.set(null); + } + } +} diff --git a/presto-jmx/src/main/java/com/facebook/presto/connector/jmx/JmxMetadata.java b/presto-jmx/src/main/java/com/facebook/presto/connector/jmx/JmxMetadata.java index f1b958959b246..a9de60e9f0699 100644 --- a/presto-jmx/src/main/java/com/facebook/presto/connector/jmx/JmxMetadata.java +++ b/presto-jmx/src/main/java/com/facebook/presto/connector/jmx/JmxMetadata.java @@ -263,11 +263,15 @@ public Map> listTableColumns(ConnectorSess } @Override - public List getTableLayouts(ConnectorSession session, ConnectorTableHandle table, Constraint constraint, Optional> desiredColumns) + public ConnectorTableLayoutResult getTableLayoutForConstraint( + ConnectorSession session, + ConnectorTableHandle table, + Constraint constraint, + Optional> desiredColumns) { JmxTableHandle handle = (JmxTableHandle) table; ConnectorTableLayout layout = new ConnectorTableLayout(new JmxTableLayoutHandle(handle, constraint.getSummary())); - return ImmutableList.of(new ConnectorTableLayoutResult(layout, constraint.getSummary())); + return new ConnectorTableLayoutResult(layout, constraint.getSummary()); } @Override diff --git a/presto-kafka/src/main/java/com/facebook/presto/kafka/KafkaMetadata.java b/presto-kafka/src/main/java/com/facebook/presto/kafka/KafkaMetadata.java index 8792b566309bc..d2320f2ebdbc6 100644 --- a/presto-kafka/src/main/java/com/facebook/presto/kafka/KafkaMetadata.java +++ b/presto-kafka/src/main/java/com/facebook/presto/kafka/KafkaMetadata.java @@ -200,7 +200,11 @@ public ColumnMetadata getColumnMetadata(ConnectorSession session, ConnectorTable } @Override - public List getTableLayouts(ConnectorSession session, ConnectorTableHandle table, Constraint constraint, Optional> desiredColumns) + public ConnectorTableLayoutResult getTableLayoutForConstraint( + ConnectorSession session, + ConnectorTableHandle table, + Constraint constraint, + Optional> desiredColumns) { KafkaTableHandle handle = convertTableHandle(table); long startTimestamp = 0; @@ -224,7 +228,7 @@ public List getTableLayouts(ConnectorSession session } ConnectorTableLayout layout = new ConnectorTableLayout(new KafkaTableLayoutHandle(handle, startTimestamp, endTimestamp)); - return ImmutableList.of(new ConnectorTableLayoutResult(layout, constraint.getSummary())); + return new ConnectorTableLayoutResult(layout, constraint.getSummary()); } @Override diff --git a/presto-kafka/src/test/java/com/facebook/presto/kafka/util/KafkaLoader.java b/presto-kafka/src/test/java/com/facebook/presto/kafka/util/KafkaLoader.java index 901baa5b23465..ddb31c6bd19cb 100644 --- a/presto-kafka/src/test/java/com/facebook/presto/kafka/util/KafkaLoader.java +++ b/presto-kafka/src/test/java/com/facebook/presto/kafka/util/KafkaLoader.java @@ -120,7 +120,7 @@ public void addResults(QueryStatusInfo statusInfo, QueryData data) } @Override - public Void build(Map setSessionProperties, Set resetSessionProperties) + public Void build(Map setSessionProperties, Set resetSessionProperties, String startTransactionId, boolean clearTransactionId) { return null; } diff --git a/presto-kudu/src/main/java/com/facebook/presto/kudu/KuduMetadata.java b/presto-kudu/src/main/java/com/facebook/presto/kudu/KuduMetadata.java index 77eccc43f7a09..603b82c56716d 100755 --- a/presto-kudu/src/main/java/com/facebook/presto/kudu/KuduMetadata.java +++ b/presto-kudu/src/main/java/com/facebook/presto/kudu/KuduMetadata.java @@ -203,7 +203,8 @@ public KuduTableHandle getTableHandle(ConnectorSession session, SchemaTableName } @Override - public List getTableLayouts(ConnectorSession session, + public ConnectorTableLayoutResult getTableLayoutForConstraint( + ConnectorSession session, ConnectorTableHandle tableHandle, Constraint constraint, Optional> desiredColumns) @@ -211,7 +212,7 @@ public List getTableLayouts(ConnectorSession session KuduTableHandle handle = (KuduTableHandle) tableHandle; ConnectorTableLayout layout = new ConnectorTableLayout( new KuduTableLayoutHandle(handle, constraint.getSummary(), desiredColumns)); - return ImmutableList.of(new ConnectorTableLayoutResult(layout, constraint.getSummary())); + return new ConnectorTableLayoutResult(layout, constraint.getSummary()); } @Override diff --git a/presto-lark-sheets/src/main/java/com/facebook/presto/lark/sheets/LarkSheetsMetadata.java b/presto-lark-sheets/src/main/java/com/facebook/presto/lark/sheets/LarkSheetsMetadata.java index 9549708c317d9..0e1ba5d8c0d6c 100644 --- a/presto-lark-sheets/src/main/java/com/facebook/presto/lark/sheets/LarkSheetsMetadata.java +++ b/presto-lark-sheets/src/main/java/com/facebook/presto/lark/sheets/LarkSheetsMetadata.java @@ -127,14 +127,15 @@ public Optional getSystemTable(ConnectorSession session, SchemaTabl } @Override - public List getTableLayouts(ConnectorSession session, + public ConnectorTableLayoutResult getTableLayoutForConstraint( + ConnectorSession session, ConnectorTableHandle table, Constraint constraint, Optional> desiredColumns) { LarkSheetsTableHandle tableHandle = (LarkSheetsTableHandle) table; ConnectorTableLayout layout = new ConnectorTableLayout(new LarkSheetsTableLayoutHandle(tableHandle)); - return ImmutableList.of(new ConnectorTableLayoutResult(layout, constraint.getSummary())); + return new ConnectorTableLayoutResult(layout, constraint.getSummary()); } @Override diff --git a/presto-local-file/src/main/java/com/facebook/presto/localfile/LocalFileMetadata.java b/presto-local-file/src/main/java/com/facebook/presto/localfile/LocalFileMetadata.java index 3db1028497578..c4224fbc7c761 100644 --- a/presto-local-file/src/main/java/com/facebook/presto/localfile/LocalFileMetadata.java +++ b/presto-local-file/src/main/java/com/facebook/presto/localfile/LocalFileMetadata.java @@ -84,11 +84,15 @@ public List listTables(ConnectorSession session, String schemaN } @Override - public List getTableLayouts(ConnectorSession session, ConnectorTableHandle table, Constraint constraint, Optional> desiredColumns) + public ConnectorTableLayoutResult getTableLayoutForConstraint( + ConnectorSession session, + ConnectorTableHandle table, + Constraint constraint, + Optional> desiredColumns) { LocalFileTableHandle tableHandle = (LocalFileTableHandle) table; ConnectorTableLayout layout = new ConnectorTableLayout(new LocalFileTableLayoutHandle(tableHandle, constraint.getSummary())); - return ImmutableList.of(new ConnectorTableLayoutResult(layout, constraint.getSummary())); + return new ConnectorTableLayoutResult(layout, constraint.getSummary()); } @Override diff --git a/presto-main-base/pom.xml b/presto-main-base/pom.xml index f724e286320a5..9c779dafd33c1 100644 --- a/presto-main-base/pom.xml +++ b/presto-main-base/pom.xml @@ -461,6 +461,20 @@ io.netty netty-transport + + + com.facebook.presto + presto-built-in-worker-function-tools + ${project.version} + test + + + + com.facebook.presto + presto-sql-invoked-functions-plugin + ${project.version} + test + @@ -525,6 +539,9 @@ com.facebook.presto.testing.assertions + + com.facebook.presto.server.MockHttpServletRequest + diff --git a/presto-main-base/src/main/java/com/facebook/presto/FullConnectorSession.java b/presto-main-base/src/main/java/com/facebook/presto/FullConnectorSession.java index 1d8ee24735d8f..d085920337a36 100644 --- a/presto-main-base/src/main/java/com/facebook/presto/FullConnectorSession.java +++ b/presto-main-base/src/main/java/com/facebook/presto/FullConnectorSession.java @@ -48,17 +48,11 @@ public class FullConnectorSession private final SessionPropertyManager sessionPropertyManager; private final SqlFunctionProperties sqlFunctionProperties; private final Map sessionFunctions; + private final RuntimeStats runtimeStats; public FullConnectorSession(Session session, ConnectorIdentity identity) { - this.session = requireNonNull(session, "session is null"); - this.identity = requireNonNull(identity, "identity is null"); - this.properties = null; - this.connectorId = null; - this.catalog = null; - this.sessionPropertyManager = null; - this.sqlFunctionProperties = session.getSqlFunctionProperties(); - this.sessionFunctions = ImmutableMap.copyOf(session.getSessionFunctions()); + this(builder(session, identity, null, null, null, null)); } public FullConnectorSession( @@ -69,14 +63,123 @@ public FullConnectorSession( String catalog, SessionPropertyManager sessionPropertyManager) { - this.session = requireNonNull(session, "session is null"); - this.identity = requireNonNull(identity, "identity is null"); - this.properties = ImmutableMap.copyOf(requireNonNull(properties, "properties is null")); - this.connectorId = requireNonNull(connectorId, "connectorId is null"); - this.catalog = requireNonNull(catalog, "catalog is null"); - this.sessionPropertyManager = requireNonNull(sessionPropertyManager, "sessionPropertyManager is null"); - this.sqlFunctionProperties = session.getSqlFunctionProperties(); - this.sessionFunctions = ImmutableMap.copyOf(session.getSessionFunctions()); + this(builder(session, identity, properties, connectorId, catalog, sessionPropertyManager)); + } + + private FullConnectorSession(Builder builder) + { + this.session = builder.getSession(); + this.identity = builder.getIdentity(); + this.properties = builder.getProperties(); + this.connectorId = builder.getConnectorId(); + this.catalog = builder.getCatalog(); + this.sessionPropertyManager = builder.getSessionPropertyManager(); + this.sqlFunctionProperties = builder.getSqlFunctionProperties() != null ? builder.getSqlFunctionProperties() : builder.getSession().getSqlFunctionProperties(); + this.sessionFunctions = builder.getSessionFunctions() != null ? builder.getSessionFunctions() : ImmutableMap.copyOf(builder.getSession().getSessionFunctions()); + this.runtimeStats = builder.getRuntimeStats() != null ? builder.getRuntimeStats() : builder.getSession().getRuntimeStats(); + } + + public static Builder builder( + Session session, + ConnectorIdentity identity, + Map properties, + ConnectorId connectorId, + String catalog, + SessionPropertyManager sessionPropertyManager) + { + return new Builder(session, identity, properties, connectorId, catalog, sessionPropertyManager); + } + + public static class Builder + { + private final Session session; + private final ConnectorIdentity identity; + private final Map properties; + private final ConnectorId connectorId; + private final String catalog; + private final SessionPropertyManager sessionPropertyManager; + + private SqlFunctionProperties sqlFunctionProperties; + private Map sessionFunctions; + private RuntimeStats runtimeStats; + + private Builder(Session session, ConnectorIdentity identity, Map properties, ConnectorId connectorId, String catalog, SessionPropertyManager sessionPropertyManager) + { + this.session = requireNonNull(session, "session is null"); + this.identity = requireNonNull(identity, "identity is null"); + this.properties = properties; + this.connectorId = connectorId; + this.catalog = catalog; + this.sessionPropertyManager = sessionPropertyManager; + } + + public Session getSession() + { + return session; + } + + public ConnectorIdentity getIdentity() + { + return identity; + } + + public Map getProperties() + { + return properties; + } + + public ConnectorId getConnectorId() + { + return connectorId; + } + + public String getCatalog() + { + return catalog; + } + + public SessionPropertyManager getSessionPropertyManager() + { + return sessionPropertyManager; + } + + public SqlFunctionProperties getSqlFunctionProperties() + { + return sqlFunctionProperties; + } + + public Builder setSqlFunctionProperties(SqlFunctionProperties sqlFunctionProperties) + { + this.sqlFunctionProperties = sqlFunctionProperties; + return this; + } + + public Map getSessionFunctions() + { + return sessionFunctions; + } + + public Builder setSessionFunctions(Map sessionFunctions) + { + this.sessionFunctions = sessionFunctions; + return this; + } + + public RuntimeStats getRuntimeStats() + { + return runtimeStats; + } + + public Builder setRuntimeStats(RuntimeStats runtimeStats) + { + this.runtimeStats = runtimeStats; + return this; + } + + public FullConnectorSession build() + { + return new FullConnectorSession(this); + } } public Session getSession() @@ -197,7 +300,7 @@ public WarningCollector getWarningCollector() @Override public RuntimeStats getRuntimeStats() { - return session.getRuntimeStats(); + return runtimeStats; } @Override diff --git a/presto-main-base/src/main/java/com/facebook/presto/Session.java b/presto-main-base/src/main/java/com/facebook/presto/Session.java index 14b7a05d142d7..05e01e17e5355 100644 --- a/presto-main-base/src/main/java/com/facebook/presto/Session.java +++ b/presto-main-base/src/main/java/com/facebook/presto/Session.java @@ -42,6 +42,7 @@ import com.facebook.presto.sql.planner.optimizations.OptimizerInformationCollector; import com.facebook.presto.sql.planner.optimizations.OptimizerResultCollector; import com.facebook.presto.transaction.TransactionManager; +import com.google.common.annotations.VisibleForTesting; import com.google.common.collect.ImmutableMap; import com.google.common.collect.ImmutableSet; import com.google.common.collect.Maps; @@ -470,6 +471,50 @@ public Session beginTransactionId(TransactionId transactionId, boolean enableRol queryType); } + @VisibleForTesting + public Session clearTransaction(TransactionManager transactionManager, AccessControl accessControl) + { + checkArgument(this.transactionId.isPresent(), "Session does not have an active transaction"); + requireNonNull(transactionManager, "transactionManager is null"); + requireNonNull(accessControl, "accessControl is null"); + + for (Entry property : systemProperties.entrySet()) { + // verify permissions + accessControl.checkCanSetSystemSessionProperty(identity, context, property.getKey()); + + // validate session property value + sessionPropertyManager.validateSystemSessionProperty(property.getKey(), property.getValue()); + } + + return new Session( + queryId, + Optional.empty(), + clientTransactionSupport, + identity, + source, + catalog, + schema, + traceToken, + timeZoneKey, + locale, + remoteUserAddress, + userAgent, + clientInfo, + clientTags, + resourceEstimates, + startTime, + systemProperties, + connectorProperties, + unprocessedCatalogProperties, + sessionPropertyManager, + preparedStatements, + sessionFunctions, + tracer, + warningCollector, + runtimeStats, + queryType); + } + public ConnectorSession toConnectorSession() { return new FullConnectorSession(this, identity.toConnectorIdentity()); @@ -495,17 +540,26 @@ public SqlFunctionProperties getSqlFunctionProperties() .build(); } - public ConnectorSession toConnectorSession(ConnectorId connectorId) + public ConnectorSession toConnectorSession(ConnectorId connectorId, RuntimeStats runtimeStats) { requireNonNull(connectorId, "connectorId is null"); - return new FullConnectorSession( - this, - identity.toConnectorIdentity(connectorId.getCatalogName()), - connectorProperties.getOrDefault(connectorId, ImmutableMap.of()), - connectorId, - connectorId.getCatalogName(), - sessionPropertyManager); + FullConnectorSession.Builder connectorSessionBuilder = FullConnectorSession + .builder( + this, + identity.toConnectorIdentity(connectorId.getCatalogName()), + connectorProperties.getOrDefault(connectorId, ImmutableMap.of()), + connectorId, + connectorId.getCatalogName(), + sessionPropertyManager) + .setRuntimeStats(runtimeStats); + + return connectorSessionBuilder.build(); + } + + public ConnectorSession toConnectorSession(ConnectorId connectorId) + { + return toConnectorSession(connectorId, runtimeStats); } public SessionRepresentation toSessionRepresentation() diff --git a/presto-main-base/src/main/java/com/facebook/presto/SystemSessionProperties.java b/presto-main-base/src/main/java/com/facebook/presto/SystemSessionProperties.java index 1a30a76798d37..f335f1f718c57 100644 --- a/presto-main-base/src/main/java/com/facebook/presto/SystemSessionProperties.java +++ b/presto-main-base/src/main/java/com/facebook/presto/SystemSessionProperties.java @@ -337,6 +337,7 @@ public final class SystemSessionProperties public static final String QUERY_CLIENT_TIMEOUT = "query_client_timeout"; public static final String REWRITE_MIN_MAX_BY_TO_TOP_N = "rewrite_min_max_by_to_top_n"; public static final String ADD_DISTINCT_BELOW_SEMI_JOIN_BUILD = "add_distinct_below_semi_join_build"; + public static final String UTILIZE_UNIQUE_PROPERTY_IN_QUERY_PLANNING = "utilize_unique_property_in_query_planning"; public static final String PUSHDOWN_SUBFIELDS_FOR_MAP_FUNCTIONS = "pushdown_subfields_for_map_functions"; public static final String MAX_SERIALIZABLE_OBJECT_SIZE = "max_serializable_object_size"; @@ -1950,6 +1951,10 @@ public SystemSessionProperties( false, value -> Duration.valueOf((String) value), Duration::toString), + booleanProperty(UTILIZE_UNIQUE_PROPERTY_IN_QUERY_PLANNING, + "Utilize the unique property of input columns in query planning", + featuresConfig.isUtilizeUniquePropertyInQueryPlanning(), + false), booleanProperty(ADD_DISTINCT_BELOW_SEMI_JOIN_BUILD, "Add distinct aggregation below semi join build", featuresConfig.isAddDistinctBelowSemiJoinBuild(), @@ -3310,6 +3315,11 @@ public static boolean isPushSubfieldsForMapFunctionsEnabled(Session session) return session.getSystemProperty(PUSHDOWN_SUBFIELDS_FOR_MAP_FUNCTIONS, Boolean.class); } + public static boolean isUtilizeUniquePropertyInQueryPlanningEnabled(Session session) + { + return session.getSystemProperty(UTILIZE_UNIQUE_PROPERTY_IN_QUERY_PLANNING, Boolean.class); + } + public static boolean isAddDistinctBelowSemiJoinBuildEnabled(Session session) { return session.getSystemProperty(ADD_DISTINCT_BELOW_SEMI_JOIN_BUILD, Boolean.class); diff --git a/presto-main-base/src/main/java/com/facebook/presto/connector/ConnectorManager.java b/presto-main-base/src/main/java/com/facebook/presto/connector/ConnectorManager.java index bbcbe517b50dc..dc85eea82f0e5 100644 --- a/presto-main-base/src/main/java/com/facebook/presto/connector/ConnectorManager.java +++ b/presto-main-base/src/main/java/com/facebook/presto/connector/ConnectorManager.java @@ -215,10 +215,10 @@ public synchronized ConnectorId createConnection(String catalogName, String conn requireNonNull(connectorName, "connectorName is null"); ConnectorFactory connectorFactory = connectorFactories.get(connectorName); checkArgument(connectorFactory != null, "No factory for connector %s", connectorName); - return createConnection(catalogName, connectorFactory, properties); + return createConnection(catalogName, connectorFactory, properties, connectorName); } - private synchronized ConnectorId createConnection(String catalogName, ConnectorFactory connectorFactory, Map properties) + private synchronized ConnectorId createConnection(String catalogName, ConnectorFactory connectorFactory, Map properties, String connectorName) { checkState(!stopped.get(), "ConnectorManager is stopped"); requireNonNull(catalogName, "catalogName is null"); @@ -229,12 +229,12 @@ private synchronized ConnectorId createConnection(String catalogName, ConnectorF ConnectorId connectorId = new ConnectorId(catalogName); checkState(!connectors.containsKey(connectorId), "A connector %s already exists", connectorId); - addCatalogConnector(catalogName, connectorId, connectorFactory, properties); + addCatalogConnector(catalogName, connectorId, connectorFactory, properties, connectorName); return connectorId; } - private synchronized void addCatalogConnector(String catalogName, ConnectorId connectorId, ConnectorFactory factory, Map properties) + private synchronized void addCatalogConnector(String catalogName, ConnectorId connectorId, ConnectorFactory factory, Map properties, String connectorName) { // create all connectors before adding, so a broken connector does not leave the system half updated MaterializedConnector connector = new MaterializedConnector(connectorId, createConnector(connectorId, factory, properties)); @@ -269,7 +269,8 @@ private synchronized void addCatalogConnector(String catalogName, ConnectorId co informationSchemaConnector.getConnectorId(), informationSchemaConnector.getConnector(), systemConnector.getConnectorId(), - systemConnector.getConnector()); + systemConnector.getConnector(), + connectorName); try { addConnectorInternal(connector); diff --git a/presto-main-base/src/main/java/com/facebook/presto/connector/system/CatalogSystemTable.java b/presto-main-base/src/main/java/com/facebook/presto/connector/system/CatalogSystemTable.java index 19cbea31463fc..e8fbcdf6a2303 100644 --- a/presto-main-base/src/main/java/com/facebook/presto/connector/system/CatalogSystemTable.java +++ b/presto-main-base/src/main/java/com/facebook/presto/connector/system/CatalogSystemTable.java @@ -15,8 +15,8 @@ import com.facebook.presto.Session; import com.facebook.presto.common.predicate.TupleDomain; +import com.facebook.presto.metadata.Catalog.CatalogContext; import com.facebook.presto.metadata.Metadata; -import com.facebook.presto.spi.ConnectorId; import com.facebook.presto.spi.ConnectorSession; import com.facebook.presto.spi.ConnectorTableMetadata; import com.facebook.presto.spi.InMemoryRecordSet; @@ -32,7 +32,7 @@ import static com.facebook.presto.common.type.VarcharType.createUnboundedVarcharType; import static com.facebook.presto.connector.system.SystemConnectorSessionUtil.toSession; -import static com.facebook.presto.metadata.MetadataListing.listCatalogs; +import static com.facebook.presto.metadata.MetadataListing.listCatalogsWithConnectorContext; import static com.facebook.presto.metadata.MetadataUtil.TableMetadataBuilder.tableMetadataBuilder; import static com.facebook.presto.spi.SystemTable.Distribution.SINGLE_COORDINATOR; import static java.util.Objects.requireNonNull; @@ -45,6 +45,7 @@ public class CatalogSystemTable public static final ConnectorTableMetadata CATALOG_TABLE = tableMetadataBuilder(CATALOG_TABLE_NAME) .column("catalog_name", createUnboundedVarcharType()) .column("connector_id", createUnboundedVarcharType()) + .column("connector_name", createUnboundedVarcharType()) .build(); private final Metadata metadata; private final AccessControl accessControl; @@ -73,8 +74,8 @@ public RecordCursor cursor(ConnectorTransactionHandle transactionHandle, Connect { Session session = toSession(transactionHandle, connectorSession); Builder table = InMemoryRecordSet.builder(CATALOG_TABLE); - for (Map.Entry entry : listCatalogs(session, metadata, accessControl).entrySet()) { - table.addRow(entry.getKey(), entry.getValue().toString()); + for (Map.Entry entry : listCatalogsWithConnectorContext(session, metadata, accessControl).entrySet()) { + table.addRow(entry.getKey(), entry.getValue().getCatalogName(), entry.getValue().getConnectorName()); } return table.build().cursor(); } diff --git a/presto-main-base/src/main/java/com/facebook/presto/event/QueryMonitor.java b/presto-main-base/src/main/java/com/facebook/presto/event/QueryMonitor.java index 8094b02e809a5..ab14cfa730639 100644 --- a/presto-main-base/src/main/java/com/facebook/presto/event/QueryMonitor.java +++ b/presto-main-base/src/main/java/com/facebook/presto/event/QueryMonitor.java @@ -50,6 +50,7 @@ import com.facebook.presto.spi.QueryId; import com.facebook.presto.spi.eventlistener.Column; import com.facebook.presto.spi.eventlistener.OperatorStatistics; +import com.facebook.presto.spi.eventlistener.OutputColumnMetadata; import com.facebook.presto.spi.eventlistener.QueryCompletedEvent; import com.facebook.presto.spi.eventlistener.QueryContext; import com.facebook.presto.spi.eventlistener.QueryCreatedEvent; @@ -600,11 +601,12 @@ private static QueryIOMetadata getQueryIOMetadata(QueryInfo queryInfo) .map(TableFinishInfo.class::cast) .findFirst(); - Optional> outputColumns = queryInfo.getOutput().get().getColumns() + Optional> outputColumnsMetadata = queryInfo.getOutput().get().getColumns() .map(columns -> columns.stream() - .map(column -> new Column( - column.getName(), - column.getType())) + .map(column -> new OutputColumnMetadata( + column.getColumnName(), + column.getColumnType(), + column.getSourceColumns())) .collect(toImmutableList())); output = Optional.of( @@ -615,7 +617,7 @@ private static QueryIOMetadata getQueryIOMetadata(QueryInfo queryInfo) tableFinishInfo.map(TableFinishInfo::getSerializedConnectorOutputMetadata), tableFinishInfo.map(TableFinishInfo::isJsonLengthLimitExceeded), queryInfo.getOutput().get().getSerializedCommitOutput(), - outputColumns)); + outputColumnsMetadata)); } return new QueryIOMetadata(inputs.build(), output); diff --git a/presto-main-base/src/main/java/com/facebook/presto/execution/AddColumnTask.java b/presto-main-base/src/main/java/com/facebook/presto/execution/AddColumnTask.java index 3fbe29ff989a2..db9f05c402538 100644 --- a/presto-main-base/src/main/java/com/facebook/presto/execution/AddColumnTask.java +++ b/presto-main-base/src/main/java/com/facebook/presto/execution/AddColumnTask.java @@ -14,7 +14,6 @@ package com.facebook.presto.execution; import com.facebook.presto.Session; -import com.facebook.presto.UnknownTypeException; import com.facebook.presto.common.QualifiedObjectName; import com.facebook.presto.common.type.Type; import com.facebook.presto.metadata.Metadata; @@ -25,6 +24,7 @@ import com.facebook.presto.spi.TableHandle; import com.facebook.presto.spi.WarningCollector; import com.facebook.presto.spi.security.AccessControl; +import com.facebook.presto.spi.type.UnknownTypeException; import com.facebook.presto.sql.analyzer.SemanticException; import com.facebook.presto.sql.tree.AddColumn; import com.facebook.presto.sql.tree.ColumnDefinition; diff --git a/presto-main-base/src/main/java/com/facebook/presto/execution/ClusterOverloadConfig.java b/presto-main-base/src/main/java/com/facebook/presto/execution/ClusterOverloadConfig.java new file mode 100644 index 0000000000000..a8ddf72707447 --- /dev/null +++ b/presto-main-base/src/main/java/com/facebook/presto/execution/ClusterOverloadConfig.java @@ -0,0 +1,111 @@ +/* + * Licensed 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 com.facebook.presto.execution; + +import com.facebook.airlift.configuration.Config; + +public class ClusterOverloadConfig +{ + public static final String OVERLOAD_POLICY_CNT_BASED = "overload_worker_cnt_based_throttling"; + public static final String OVERLOAD_POLICY_PCT_BASED = "overload_worker_pct_based_throttling"; + private boolean clusterOverloadThrottlingEnabled; + private double allowedOverloadWorkersPct = 0.01; + private int allowedOverloadWorkersCnt; + private String overloadPolicyType = OVERLOAD_POLICY_CNT_BASED; + private int overloadCheckCacheTtlInSecs = 5; + + /** + * Gets the time-to-live for the cached cluster overload state. + * This determines how frequently the system will re-evaluate whether the cluster is overloaded. + * + * @return the cache TTL duration + */ + public int getOverloadCheckCacheTtlInSecs() + { + return overloadCheckCacheTtlInSecs; + } + + /** + * Gets the time-to-live for the cached cluster overload state. + * This determines how frequently the system will re-evaluate whether the cluster is overloaded. + * + * @return the cache TTL duration + */ + public int getOverloadCheckCacheTtlMillis() + { + return overloadCheckCacheTtlInSecs * 1000; + } + + /** + * Sets the time-to-live for the cached cluster overload state. + * + * @param overloadCheckCacheTtlInSecs the cache TTL duration + * @return this for chaining + */ + @Config("cluster.overload-check-cache-ttl-secs") + public ClusterOverloadConfig setOverloadCheckCacheTtlInSecs(int overloadCheckCacheTtlInSecs) + { + this.overloadCheckCacheTtlInSecs = overloadCheckCacheTtlInSecs; + return this; + } + + @Config("cluster-overload.enable-throttling") + public ClusterOverloadConfig setClusterOverloadThrottlingEnabled(boolean clusterOverloadThrottlingEnabled) + { + this.clusterOverloadThrottlingEnabled = clusterOverloadThrottlingEnabled; + return this; + } + + public boolean isClusterOverloadThrottlingEnabled() + { + return this.clusterOverloadThrottlingEnabled; + } + + @Config("cluster-overload.allowed-overload-workers-pct") + public ClusterOverloadConfig setAllowedOverloadWorkersPct(Double allowedOverloadWorkersPct) + { + this.allowedOverloadWorkersPct = allowedOverloadWorkersPct; + return this; + } + + public double getAllowedOverloadWorkersPct() + { + return this.allowedOverloadWorkersPct; + } + + @Config("cluster-overload.allowed-overload-workers-cnt") + public ClusterOverloadConfig setAllowedOverloadWorkersCnt(int allowedOverloadWorkersCnt) + { + this.allowedOverloadWorkersCnt = allowedOverloadWorkersCnt; + return this; + } + + public double getAllowedOverloadWorkersCnt() + { + return this.allowedOverloadWorkersCnt; + } + + @Config("cluster-overload.overload-policy-type") + public ClusterOverloadConfig setOverloadPolicyType(String overloadPolicyType) + { + // validate + this.overloadPolicyType = overloadPolicyType; + return this; + } + + public String getOverloadPolicyType() + { + return this.overloadPolicyType; + } +} diff --git a/presto-main-base/src/main/java/com/facebook/presto/execution/CreateTableTask.java b/presto-main-base/src/main/java/com/facebook/presto/execution/CreateTableTask.java index f2199283cf43b..306bbde275f6e 100644 --- a/presto-main-base/src/main/java/com/facebook/presto/execution/CreateTableTask.java +++ b/presto-main-base/src/main/java/com/facebook/presto/execution/CreateTableTask.java @@ -14,7 +14,6 @@ package com.facebook.presto.execution; import com.facebook.presto.Session; -import com.facebook.presto.UnknownTypeException; import com.facebook.presto.common.QualifiedObjectName; import com.facebook.presto.common.type.Type; import com.facebook.presto.metadata.Metadata; @@ -28,6 +27,7 @@ import com.facebook.presto.spi.constraints.PrimaryKeyConstraint; import com.facebook.presto.spi.constraints.TableConstraint; import com.facebook.presto.spi.security.AccessControl; +import com.facebook.presto.spi.type.UnknownTypeException; import com.facebook.presto.sql.analyzer.SemanticException; import com.facebook.presto.sql.tree.ColumnDefinition; import com.facebook.presto.sql.tree.ConstraintSpecification; diff --git a/presto-main-base/src/main/java/com/facebook/presto/execution/Output.java b/presto-main-base/src/main/java/com/facebook/presto/execution/Output.java index 655e427b53a27..924bb4243d1d5 100644 --- a/presto-main-base/src/main/java/com/facebook/presto/execution/Output.java +++ b/presto-main-base/src/main/java/com/facebook/presto/execution/Output.java @@ -14,6 +14,7 @@ package com.facebook.presto.execution; import com.facebook.presto.spi.ConnectorId; +import com.facebook.presto.spi.eventlistener.OutputColumnMetadata; import com.fasterxml.jackson.annotation.JsonCreator; import com.fasterxml.jackson.annotation.JsonProperty; import com.google.common.collect.ImmutableList; @@ -32,7 +33,7 @@ public final class Output private final String schema; private final String table; private final String serializedCommitOutput; - private final Optional> columns; + private final Optional> columns; @JsonCreator public Output( @@ -40,7 +41,7 @@ public Output( @JsonProperty("schema") String schema, @JsonProperty("table") String table, @JsonProperty("serializedCommitOutput") String serializedCommitOutput, - @JsonProperty("columns") Optional> columns) + @JsonProperty("columns") Optional> columns) { this.connectorId = requireNonNull(connectorId, "connectorId is null"); this.schema = requireNonNull(schema, "schema is null"); @@ -74,7 +75,7 @@ public String getSerializedCommitOutput() } @JsonProperty - public Optional> getColumns() + public Optional> getColumns() { return columns; } diff --git a/presto-main-base/src/main/java/com/facebook/presto/execution/resourceGroups/InternalResourceGroup.java b/presto-main-base/src/main/java/com/facebook/presto/execution/resourceGroups/InternalResourceGroup.java index e1504e7b50f21..82274c04b3550 100644 --- a/presto-main-base/src/main/java/com/facebook/presto/execution/resourceGroups/InternalResourceGroup.java +++ b/presto-main-base/src/main/java/com/facebook/presto/execution/resourceGroups/InternalResourceGroup.java @@ -18,6 +18,7 @@ import com.facebook.airlift.units.Duration; import com.facebook.presto.execution.ManagedQueryExecution; import com.facebook.presto.execution.resourceGroups.WeightedFairQueue.Usage; +import com.facebook.presto.execution.scheduler.clusterOverload.ClusterResourceChecker; import com.facebook.presto.metadata.InternalNodeManager; import com.facebook.presto.server.QueryStateInfo; import com.facebook.presto.server.ResourceGroupInfo; @@ -96,6 +97,7 @@ public class InternalResourceGroup private final Function> additionalRuntimeInfo; private final Predicate shouldWaitForResourceManagerUpdate; private final InternalNodeManager nodeManager; + private final ClusterResourceChecker clusterResourceChecker; // Configuration // ============= @@ -166,12 +168,14 @@ protected InternalResourceGroup( boolean staticResourceGroup, Function> additionalRuntimeInfo, Predicate shouldWaitForResourceManagerUpdate, - InternalNodeManager nodeManager) + InternalNodeManager nodeManager, + ClusterResourceChecker clusterResourceChecker) { this.parent = requireNonNull(parent, "parent is null"); this.jmxExportListener = requireNonNull(jmxExportListener, "jmxExportListener is null"); this.executor = requireNonNull(executor, "executor is null"); this.nodeManager = requireNonNull(nodeManager, "node manager is null"); + this.clusterResourceChecker = requireNonNull(clusterResourceChecker, "clusterResourceChecker is null"); requireNonNull(name, "name is null"); if (parent.isPresent()) { id = new ResourceGroupId(parent.get().id, name); @@ -671,7 +675,8 @@ public InternalResourceGroup getOrCreateSubGroup(String name, boolean staticSegm staticResourceGroup && staticSegment, additionalRuntimeInfo, shouldWaitForResourceManagerUpdate, - nodeManager); + nodeManager, + clusterResourceChecker); // Sub group must use query priority to ensure ordering if (schedulingPolicy == QUERY_PRIORITY) { subGroup.setSchedulingPolicy(QUERY_PRIORITY); @@ -770,7 +775,7 @@ private void enqueueQuery(ManagedQueryExecution query) } // This method must be called whenever the group's eligibility to run more queries may have changed. - private void updateEligibility() + protected void updateEligibility() { checkState(Thread.holdsLock(root), "Must hold lock to update eligibility"); synchronized (root) { @@ -1019,6 +1024,11 @@ private boolean canRunMore() { checkState(Thread.holdsLock(root), "Must hold lock"); synchronized (root) { + // Check if more queries can be run on the cluster based on cluster overload + if (clusterResourceChecker.isClusterCurrentlyOverloaded()) { + return false; + } + if (cpuUsageMillis >= hardCpuLimitMillis) { return false; } @@ -1135,7 +1145,8 @@ public RootInternalResourceGroup( Executor executor, Function> additionalRuntimeInfo, Predicate shouldWaitForResourceManagerUpdate, - InternalNodeManager nodeManager) + InternalNodeManager nodeManager, + ClusterResourceChecker clusterResourceChecker) { super(Optional.empty(), name, @@ -1144,7 +1155,16 @@ public RootInternalResourceGroup( true, additionalRuntimeInfo, shouldWaitForResourceManagerUpdate, - nodeManager); + nodeManager, + clusterResourceChecker); + } + + public synchronized void updateEligibilityRecursively(InternalResourceGroup group) + { + group.updateEligibility(); + for (InternalResourceGroup subGroup : group.subGroups()) { + updateEligibilityRecursively(subGroup); + } } public synchronized void processQueuedQueries() diff --git a/presto-main-base/src/main/java/com/facebook/presto/execution/resourceGroups/InternalResourceGroupManager.java b/presto-main-base/src/main/java/com/facebook/presto/execution/resourceGroups/InternalResourceGroupManager.java index 6dbf4115d0442..f760cfa19b031 100644 --- a/presto-main-base/src/main/java/com/facebook/presto/execution/resourceGroups/InternalResourceGroupManager.java +++ b/presto-main-base/src/main/java/com/facebook/presto/execution/resourceGroups/InternalResourceGroupManager.java @@ -19,6 +19,8 @@ import com.facebook.presto.execution.ManagedQueryExecution; import com.facebook.presto.execution.QueryManagerConfig; import com.facebook.presto.execution.resourceGroups.InternalResourceGroup.RootInternalResourceGroup; +import com.facebook.presto.execution.scheduler.clusterOverload.ClusterOverloadStateListener; +import com.facebook.presto.execution.scheduler.clusterOverload.ClusterResourceChecker; import com.facebook.presto.metadata.InternalNodeManager; import com.facebook.presto.resourcemanager.ResourceGroupService; import com.facebook.presto.server.ResourceGroupInfo; @@ -81,7 +83,7 @@ @ThreadSafe public final class InternalResourceGroupManager - implements ResourceGroupManager + implements ResourceGroupManager, ClusterOverloadStateListener { private static final Logger log = Logger.get(InternalResourceGroupManager.class); private static final File RESOURCE_GROUPS_CONFIGURATION = new File("etc/resource-groups.properties"); @@ -112,6 +114,7 @@ public final class InternalResourceGroupManager private final QueryManagerConfig queryManagerConfig; private final InternalNodeManager nodeManager; private AtomicBoolean isConfigurationManagerLoaded; + private final ClusterResourceChecker clusterResourceChecker; @Inject public InternalResourceGroupManager( @@ -121,7 +124,8 @@ public InternalResourceGroupManager( MBeanExporter exporter, ResourceGroupService resourceGroupService, ServerConfig serverConfig, - InternalNodeManager nodeManager) + InternalNodeManager nodeManager, + ClusterResourceChecker clusterResourceChecker) { this.queryManagerConfig = requireNonNull(queryManagerConfig, "queryManagerConfig is null"); this.exporter = requireNonNull(exporter, "exporter is null"); @@ -137,6 +141,7 @@ public InternalResourceGroupManager( this.resourceGroupRuntimeExecutor = new PeriodicTaskExecutor(resourceGroupRuntimeInfoRefreshInterval.toMillis(), refreshExecutor, this::refreshResourceGroupRuntimeInfo); configurationManagerFactories.putIfAbsent(LegacyResourceGroupConfigurationManager.NAME, new LegacyResourceGroupConfigurationManager.Factory()); this.isConfigurationManagerLoaded = new AtomicBoolean(false); + this.clusterResourceChecker = clusterResourceChecker; } @Override @@ -254,6 +259,8 @@ public ResourceGroupConfigurationManager getConfigurationManager() @PreDestroy public void destroy() { + // Unregister from cluster overload state changes + clusterResourceChecker.removeListener(this); refreshExecutor.shutdownNow(); resourceGroupRuntimeExecutor.stop(); } @@ -275,6 +282,9 @@ public void start() if (isResourceManagerEnabled) { resourceGroupRuntimeExecutor.start(); } + + // Register as listener for cluster overload state changes + clusterResourceChecker.addListener(this); } } @@ -396,7 +406,7 @@ private synchronized void createGroupIfNecessary(SelectionContext context, Ex else { RootInternalResourceGroup root; if (!isResourceManagerEnabled) { - root = new RootInternalResourceGroup(id.getSegments().get(0), this::exportGroup, executor, ignored -> Optional.empty(), rg -> false, nodeManager); + root = new RootInternalResourceGroup(id.getSegments().get(0), this::exportGroup, executor, ignored -> Optional.empty(), rg -> false, nodeManager, clusterResourceChecker); } else { root = new RootInternalResourceGroup( @@ -409,7 +419,8 @@ private synchronized void createGroupIfNecessary(SelectionContext context, Ex resourceGroupRuntimeInfosSnapshot::get, lastUpdatedResourceGroupRuntimeInfo::get, concurrencyThreshold), - nodeManager); + nodeManager, + clusterResourceChecker); } group = root; rootGroups.add(root); @@ -463,6 +474,24 @@ public int getQueriesQueuedOnInternal() return queriesQueuedInternal; } + @Override + public void onClusterEnteredOverloadedState() + { + // Resource groups will handle overload state through their existing admission control logic + // No additional action needed here as queries will be queued automatically + } + + @Override + public void onClusterExitedOverloadedState() + { + log.info("Cluster exited overloaded state, updating eligibility for all resource groups"); + for (RootInternalResourceGroup rootGroup : rootGroups) { + synchronized (rootGroup) { + rootGroup.updateEligibilityRecursively(rootGroup); + } + } + } + @Managed public long getLastSchedulingCycleRuntimeDelayMs() { diff --git a/presto-main-base/src/main/java/com/facebook/presto/execution/scheduler/clusterOverload/ClusterOverloadPolicy.java b/presto-main-base/src/main/java/com/facebook/presto/execution/scheduler/clusterOverload/ClusterOverloadPolicy.java new file mode 100644 index 0000000000000..175c18afe5405 --- /dev/null +++ b/presto-main-base/src/main/java/com/facebook/presto/execution/scheduler/clusterOverload/ClusterOverloadPolicy.java @@ -0,0 +1,40 @@ +/* + * Licensed 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 com.facebook.presto.execution.scheduler.clusterOverload; + +import com.facebook.presto.metadata.InternalNodeManager; + +/** + * Interface for policies that determine if cluster is overloaded. + * Implementations can check various metrics from NodeStats to determine + * if a worker is overloaded and queries should be throttled. + */ +public interface ClusterOverloadPolicy +{ + /** + * Checks if cluster is overloaded. + * + * @param nodeManager The node manager to get node information + * @return true if cluster is overloaded, false otherwise + */ + boolean isClusterOverloaded(InternalNodeManager nodeManager); + + /** + * Gets the name of the policy. + * + * @return The name of the policy + */ + String getName(); +} diff --git a/presto-main-base/src/main/java/com/facebook/presto/execution/scheduler/clusterOverload/ClusterOverloadPolicyFactory.java b/presto-main-base/src/main/java/com/facebook/presto/execution/scheduler/clusterOverload/ClusterOverloadPolicyFactory.java new file mode 100644 index 0000000000000..2d6b9b82a8962 --- /dev/null +++ b/presto-main-base/src/main/java/com/facebook/presto/execution/scheduler/clusterOverload/ClusterOverloadPolicyFactory.java @@ -0,0 +1,78 @@ +/* + * Licensed 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 com.facebook.presto.execution.scheduler.clusterOverload; + +import com.google.common.collect.ImmutableMap; +import jakarta.inject.Inject; + +import java.util.Map; +import java.util.Optional; + +import static java.util.Objects.requireNonNull; + +/** + * Factory for creating ClusterOverloadPolicy instances. + * This allows for extensible policy creation based on configuration. + */ +public class ClusterOverloadPolicyFactory +{ + private final Map policies; + + @Inject + public ClusterOverloadPolicyFactory(ClusterOverloadPolicy clusterOverloadPolicy) + { + requireNonNull(clusterOverloadPolicy, "clusterOverloadPolicy is null"); + + // Register available policies + ImmutableMap.Builder policiesBuilder = ImmutableMap.builder(); + + // Add the default overload policy - use the injected instance + policiesBuilder.put(clusterOverloadPolicy.getName(), clusterOverloadPolicy); + + // Add more policies here as they are implemented + this.policies = policiesBuilder.build(); + } + + /** + * Get a policy by name. + * + * @param name The name of the policy to get + * @return The policy, or empty if no policy with that name exists + */ + public Optional getPolicy(String name) + { + return Optional.ofNullable(policies.get(name)); + } + + /** + * Get the default policy. + * + * @return The default policy + */ + public ClusterOverloadPolicy getDefaultPolicy() + { + // Default to CPU/Memory policy + return policies.get("cpu-memory-overload"); + } + + /** + * Get all available policies. + * + * @return Map of policy name to policy + */ + public Map getPolicies() + { + return policies; + } +} diff --git a/presto-main-base/src/main/java/com/facebook/presto/execution/scheduler/clusterOverload/ClusterOverloadPolicyModule.java b/presto-main-base/src/main/java/com/facebook/presto/execution/scheduler/clusterOverload/ClusterOverloadPolicyModule.java new file mode 100644 index 0000000000000..207a0fd6dd198 --- /dev/null +++ b/presto-main-base/src/main/java/com/facebook/presto/execution/scheduler/clusterOverload/ClusterOverloadPolicyModule.java @@ -0,0 +1,44 @@ +/* + * Licensed 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 com.facebook.presto.execution.scheduler.clusterOverload; + +import com.facebook.airlift.configuration.AbstractConfigurationAwareModule; +import com.facebook.presto.execution.ClusterOverloadConfig; +import com.google.inject.Binder; + +import static com.facebook.airlift.configuration.ConfigBinder.configBinder; +import static com.google.inject.Scopes.SINGLETON; + +/** + * Provides bindings for the node overload policy and cluster resource checker. + */ +public class ClusterOverloadPolicyModule + extends AbstractConfigurationAwareModule +{ + @Override + protected void setup(Binder binder) + { + // Bind the default node overload policy + binder.bind(ClusterOverloadPolicy.class).to(CpuMemoryOverloadPolicy.class).in(SINGLETON); + + // Bind the node overload policy factory + binder.bind(ClusterOverloadPolicyFactory.class).in(SINGLETON); + + // Bind the cluster resource checker + binder.bind(ClusterResourceChecker.class).in(SINGLETON); + + // Bind the cluster overload config + configBinder(binder).bindConfig(ClusterOverloadConfig.class); + } +} diff --git a/presto-main-base/src/main/java/com/facebook/presto/execution/scheduler/clusterOverload/ClusterOverloadStateListener.java b/presto-main-base/src/main/java/com/facebook/presto/execution/scheduler/clusterOverload/ClusterOverloadStateListener.java new file mode 100644 index 0000000000000..86d0845545e01 --- /dev/null +++ b/presto-main-base/src/main/java/com/facebook/presto/execution/scheduler/clusterOverload/ClusterOverloadStateListener.java @@ -0,0 +1,37 @@ +/* + * Licensed 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 com.facebook.presto.execution.scheduler.clusterOverload; + +/** + * Listener interface for receiving notifications about cluster overload state changes. + * This interface allows components to react to changes in cluster capacity and + * resource availability, particularly when the cluster transitions from an overloaded + * state back to a normal state. + */ +public interface ClusterOverloadStateListener +{ + /** + * Called when the cluster enters an overloaded state. + * This indicates that the cluster resources are under stress and + * new query admissions should be restricted. + */ + void onClusterEnteredOverloadedState(); + + /** + * Called when the cluster exits the overloaded state. + * This indicates that cluster resources have recovered and + * normal query processing can resume. + */ + void onClusterExitedOverloadedState(); +} diff --git a/presto-main-base/src/main/java/com/facebook/presto/execution/scheduler/clusterOverload/ClusterResourceChecker.java b/presto-main-base/src/main/java/com/facebook/presto/execution/scheduler/clusterOverload/ClusterResourceChecker.java new file mode 100644 index 0000000000000..683fd703e970d --- /dev/null +++ b/presto-main-base/src/main/java/com/facebook/presto/execution/scheduler/clusterOverload/ClusterResourceChecker.java @@ -0,0 +1,248 @@ +/* + * Licensed 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 com.facebook.presto.execution.scheduler.clusterOverload; + +import com.facebook.airlift.log.Logger; +import com.facebook.airlift.stats.CounterStat; +import com.facebook.airlift.stats.TimeStat; +import com.facebook.presto.execution.ClusterOverloadConfig; +import com.facebook.presto.metadata.InternalNodeManager; +import com.google.errorprone.annotations.ThreadSafe; +import com.google.inject.Singleton; +import jakarta.annotation.PostConstruct; +import jakarta.annotation.PreDestroy; +import jakarta.inject.Inject; +import org.weakref.jmx.Managed; +import org.weakref.jmx.Nested; + +import java.util.concurrent.CopyOnWriteArrayList; +import java.util.concurrent.ScheduledExecutorService; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicBoolean; +import java.util.concurrent.atomic.AtomicLong; + +import static com.facebook.airlift.concurrent.Threads.threadsNamed; +import static java.util.Objects.requireNonNull; +import static java.util.concurrent.Executors.newSingleThreadScheduledExecutor; + +/** + * Provides methods to check if more queries can be run on the cluster + * based on various resource constraints. + */ +@Singleton +@ThreadSafe +public class ClusterResourceChecker +{ + private static final Logger log = Logger.get(ClusterResourceChecker.class); + + private final ClusterOverloadPolicy clusterOverloadPolicy; + private final ClusterOverloadConfig config; + private final AtomicBoolean cachedOverloadState = new AtomicBoolean(false); + private final AtomicLong lastCheckTimeMillis = new AtomicLong(0); + private final CounterStat overloadDetectionCount = new CounterStat(); + private final TimeStat timeSinceLastCheck = new TimeStat(); + private final AtomicLong overloadStartTimeMillis = new AtomicLong(0); + private final CopyOnWriteArrayList listeners = new CopyOnWriteArrayList<>(); + + private final ScheduledExecutorService overloadCheckerExecutor; + private final InternalNodeManager nodeManager; + + @Inject + public ClusterResourceChecker(ClusterOverloadPolicy clusterOverloadPolicy, ClusterOverloadConfig config, InternalNodeManager nodeManager) + { + this.clusterOverloadPolicy = requireNonNull(clusterOverloadPolicy, "clusterOverloadPolicy is null"); + this.config = requireNonNull(config, "config is null"); + this.nodeManager = requireNonNull(nodeManager, "nodeManager is null"); + this.overloadCheckerExecutor = newSingleThreadScheduledExecutor(threadsNamed("cluster-overload-checker-%s")); + } + + @PostConstruct + public void start() + { + if (config.isClusterOverloadThrottlingEnabled()) { + long checkIntervalMillis = Math.max(1000, config.getOverloadCheckCacheTtlInSecs() * 1000L); + overloadCheckerExecutor.scheduleWithFixedDelay(() -> { + try { + performPeriodicOverloadCheck(); + } + catch (Exception e) { + log.error(e, "Error polling cluster overload state"); + } + }, checkIntervalMillis, checkIntervalMillis, TimeUnit.MILLISECONDS); + log.info("Started periodic cluster overload checker with interval: %d milliseconds", checkIntervalMillis); + // Perform initial check + performPeriodicOverloadCheck(); + } + } + + @PreDestroy + public void stop() + { + overloadCheckerExecutor.shutdownNow(); + } + + /** + * Registers a listener to be notified when the cluster exits the overloaded state. + * + * @param listener the listener to register + */ + public void addListener(ClusterOverloadStateListener listener) + { + requireNonNull(listener, "listener is null"); + listeners.add(listener); + } + + /** + * Removes a previously registered listener. + * + * @param listener the listener to remove + */ + public void removeListener(ClusterOverloadStateListener listener) + { + listeners.remove(listener); + } + + /** + * Performs a periodic check of cluster overload state. + * This method is called by the periodic task when throttling is enabled. + * Updates JMX metrics and notifies listeners when cluster exits overloaded state. + */ + private void performPeriodicOverloadCheck() + { + try { + long currentTimeMillis = System.currentTimeMillis(); + long lastCheckTime = lastCheckTimeMillis.get(); + + if (lastCheckTime > 0) { + timeSinceLastCheck.add(currentTimeMillis - lastCheckTime, TimeUnit.MILLISECONDS); + } + + boolean isOverloaded = clusterOverloadPolicy.isClusterOverloaded(nodeManager); + synchronized (this) { + boolean wasOverloaded = cachedOverloadState.getAndSet(isOverloaded); + lastCheckTimeMillis.set(currentTimeMillis); + + if (isOverloaded && !wasOverloaded) { + overloadDetectionCount.update(1); + overloadStartTimeMillis.set(currentTimeMillis); + log.info("Cluster entered overloaded state via periodic check"); + } + else if (!isOverloaded && wasOverloaded) { + long overloadDuration = currentTimeMillis - overloadStartTimeMillis.get(); + log.info("Cluster exited overloaded state after %d ms via periodic check", overloadDuration); + overloadStartTimeMillis.set(0); + // Notify listeners that cluster exited overload state + notifyClusterExitedOverloadedState(); + } + } + + log.debug("Periodic overload check completed: %s", isOverloaded ? "OVERLOADED" : "NOT OVERLOADED"); + } + catch (Exception e) { + log.error(e, "Error during periodic cluster overload check"); + } + } + + /** + * Returns the current overload state of the cluster. + * @return true if cluster is overloaded, false otherwise + */ + public boolean isClusterCurrentlyOverloaded() + { + if (!config.isClusterOverloadThrottlingEnabled()) { + return false; + } + + return cachedOverloadState.get(); + } + + /** + * Notifies all registered listeners that the cluster has exited the overloaded state. + */ + private void notifyClusterExitedOverloadedState() + { + for (ClusterOverloadStateListener listener : listeners) { + listener.onClusterExitedOverloadedState(); + } + } + + /** + * Returns whether cluster overload throttling is enabled. + * When disabled, the cluster overload check will be bypassed. + * + * @return true if throttling is enabled, false otherwise + */ + @Managed + public boolean isClusterOverloadThrottlingEnabled() + { + return config.isClusterOverloadThrottlingEnabled(); + } + + /** + * Returns whether the cluster is currently in an overloaded state. + * This is exposed as a JMX metric for monitoring. + * + * @return true if the cluster is overloaded, false otherwise + */ + @Managed + public boolean isClusterOverloaded() + { + return cachedOverloadState.get(); + } + + /** + * Returns the number of times the cluster has entered an overloaded state. + * + * @return counter of overload detections + */ + @Managed + @Nested + public CounterStat getOverloadDetectionCount() + { + return overloadDetectionCount; + } + + /** + * Returns statistics about the time between overload checks. + * + * @return time statistics for overload checks + */ + @Managed + @Nested + public TimeStat getTimeSinceLastCheck() + { + return timeSinceLastCheck; + } + + /** + * Returns the duration in milliseconds that the cluster has been in an overloaded state. + * Returns 0 if the cluster is not currently overloaded. + * + * Note: This method reads two atomic fields but doesn't need synchronization because: + * 1. Single writer (periodic task) ensures consistent updates + * 2. Atomic fields provide memory visibility guarantees + * 3. Slight inconsistency in edge cases is acceptable for monitoring metrics + * + * @return duration in milliseconds of current overload state + */ + @Managed + public long getOverloadDurationMillis() + { + long startTime = overloadStartTimeMillis.get(); + if (startTime == 0 || !cachedOverloadState.get()) { + return 0; + } + return System.currentTimeMillis() - startTime; + } +} diff --git a/presto-main-base/src/main/java/com/facebook/presto/execution/scheduler/clusterOverload/CpuMemoryOverloadPolicy.java b/presto-main-base/src/main/java/com/facebook/presto/execution/scheduler/clusterOverload/CpuMemoryOverloadPolicy.java new file mode 100644 index 0000000000000..14d5387454178 --- /dev/null +++ b/presto-main-base/src/main/java/com/facebook/presto/execution/scheduler/clusterOverload/CpuMemoryOverloadPolicy.java @@ -0,0 +1,169 @@ +/* + * Licensed 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 com.facebook.presto.execution.scheduler.clusterOverload; + +import com.facebook.airlift.log.Logger; +import com.facebook.presto.execution.ClusterOverloadConfig; +import com.facebook.presto.metadata.InternalNode; +import com.facebook.presto.metadata.InternalNodeManager; +import com.facebook.presto.spi.NodeLoadMetrics; +import jakarta.inject.Inject; + +import java.util.HashSet; +import java.util.Optional; +import java.util.Set; + +import static com.facebook.presto.spi.NodeState.ACTIVE; +import static java.lang.String.format; +import static java.util.Objects.requireNonNull; + +/** + * A policy that checks if cluster is overloaded based on CPU or memory metrics. + * Supports two modes of operation: + * - Percentage-based: Checks if the percentage of overloaded workers exceeds a threshold + * - Count-based: Checks if the absolute count of overloaded workers exceeds a threshold + */ +public class CpuMemoryOverloadPolicy + implements ClusterOverloadPolicy +{ + private static final Logger log = Logger.get(CpuMemoryOverloadPolicy.class); + + private final double allowedOverloadWorkersPct; + private final double allowedOverloadWorkersCnt; + private final String policyType; + + @Inject + public CpuMemoryOverloadPolicy(ClusterOverloadConfig config) + { + this.allowedOverloadWorkersPct = config.getAllowedOverloadWorkersPct(); + this.allowedOverloadWorkersCnt = config.getAllowedOverloadWorkersCnt(); + this.policyType = requireNonNull(config.getOverloadPolicyType(), "policyType is null"); + } + + @Override + public boolean isClusterOverloaded(InternalNodeManager nodeManager) + { + Set activeNodes = nodeManager.getNodes(ACTIVE); + if (activeNodes.isEmpty()) { + return false; + } + + OverloadStats stats = collectOverloadStats(activeNodes, nodeManager); + return evaluateOverload(stats, activeNodes.size()); + } + + private OverloadStats collectOverloadStats(Set activeNodes, InternalNodeManager nodeManager) + { + Set overloadedNodeIds = new HashSet<>(); + int overloadedWorkersCnt = 0; + + for (InternalNode node : activeNodes) { + Optional metricsOptional = nodeManager.getNodeLoadMetrics(node.getNodeIdentifier()); + String nodeId = node.getNodeIdentifier(); + + if (!metricsOptional.isPresent()) { + continue; + } + + NodeLoadMetrics metrics = metricsOptional.get(); + + // Check for CPU overload + if (metrics.getCpuOverload()) { + overloadedNodeIds.add(String.format("%s (CPU overloaded)", nodeId)); + overloadedWorkersCnt++; + continue; + } + + // Check for memory overload + if (metrics.getMemoryOverload()) { + overloadedNodeIds.add(String.format("%s (Memory overloaded)", nodeId)); + overloadedWorkersCnt++; + } + } + + return new OverloadStats(overloadedNodeIds, overloadedWorkersCnt); + } + + private boolean evaluateOverload(OverloadStats stats, int totalNodes) + { + boolean isOverloaded; + + if (ClusterOverloadConfig.OVERLOAD_POLICY_PCT_BASED.equals(policyType)) { + double overloadedWorkersPct = (double) stats.getOverloadedWorkersCnt() / totalNodes; + isOverloaded = overloadedWorkersPct > allowedOverloadWorkersPct; + + if (isOverloaded) { + logOverload( + String.format("%s%% of workers are overloaded (threshold: %s%%)", + format("%.2f", overloadedWorkersPct * 100), + format("%.2f", allowedOverloadWorkersPct * 100)), + stats.getOverloadedNodeIds()); + } + } + else if (ClusterOverloadConfig.OVERLOAD_POLICY_CNT_BASED.equals(policyType)) { + isOverloaded = stats.getOverloadedWorkersCnt() > allowedOverloadWorkersCnt; + + if (isOverloaded) { + logOverload( + String.format("%s workers are overloaded (threshold: %s workers)", + stats.getOverloadedWorkersCnt(), allowedOverloadWorkersCnt), + stats.getOverloadedNodeIds()); + } + } + else { + throw new IllegalStateException("Unknown cluster overload policy type: " + policyType); + } + + return isOverloaded; + } + + private void logOverload(String message, Set overloadedNodeIds) + { + log.warn("Cluster is overloaded: " + message); + if (!overloadedNodeIds.isEmpty()) { + log.warn("Overloaded nodes: %s", String.join(", ", overloadedNodeIds)); + } + } + + @Override + public String getName() + { + return "cpu-memory-overload-" + + (ClusterOverloadConfig.OVERLOAD_POLICY_PCT_BASED.equals(policyType) ? "pct" : "cnt"); + } + + // Helper class to encapsulate overload statistics + private static class OverloadStats + { + private final Set overloadedNodeIds; + private final int overloadedWorkersCnt; + + public OverloadStats(Set overloadedNodeIds, int overloadedWorkersCnt) + { + this.overloadedNodeIds = overloadedNodeIds; + this.overloadedWorkersCnt = overloadedWorkersCnt; + } + + public Set getOverloadedNodeIds() + { + return overloadedNodeIds; + } + + public int getOverloadedWorkersCnt() + { + return overloadedWorkersCnt; + } + } +} diff --git a/presto-main-base/src/main/java/com/facebook/presto/metadata/BuiltInFunctionKind.java b/presto-main-base/src/main/java/com/facebook/presto/metadata/BuiltInFunctionKind.java index 4d12bb7f97d16..71ca25dd468cc 100644 --- a/presto-main-base/src/main/java/com/facebook/presto/metadata/BuiltInFunctionKind.java +++ b/presto-main-base/src/main/java/com/facebook/presto/metadata/BuiltInFunctionKind.java @@ -20,7 +20,8 @@ public enum BuiltInFunctionKind { ENGINE(0), - PLUGIN(1); + PLUGIN(1), + WORKER(2); private final int value; diff --git a/presto-main-base/src/main/java/com/facebook/presto/metadata/BuiltInPluginFunctionNamespaceManager.java b/presto-main-base/src/main/java/com/facebook/presto/metadata/BuiltInPluginFunctionNamespaceManager.java index 6a694f436f17e..e0cf7dc25472f 100644 --- a/presto-main-base/src/main/java/com/facebook/presto/metadata/BuiltInPluginFunctionNamespaceManager.java +++ b/presto-main-base/src/main/java/com/facebook/presto/metadata/BuiltInPluginFunctionNamespaceManager.java @@ -13,226 +13,38 @@ */ package com.facebook.presto.metadata; -import com.facebook.presto.common.Page; -import com.facebook.presto.common.QualifiedObjectName; -import com.facebook.presto.common.block.BlockEncodingSerde; -import com.facebook.presto.common.function.SqlFunctionResult; -import com.facebook.presto.common.type.TypeManager; -import com.facebook.presto.common.type.UserDefinedType; -import com.facebook.presto.spi.PrestoException; -import com.facebook.presto.spi.function.AlterRoutineCharacteristics; -import com.facebook.presto.spi.function.FunctionHandle; -import com.facebook.presto.spi.function.FunctionMetadata; -import com.facebook.presto.spi.function.FunctionNamespaceManager; -import com.facebook.presto.spi.function.FunctionNamespaceTransactionHandle; -import com.facebook.presto.spi.function.Parameter; -import com.facebook.presto.spi.function.ScalarFunctionImplementation; -import com.facebook.presto.spi.function.Signature; +import com.facebook.presto.spi.function.FunctionImplementationType; import com.facebook.presto.spi.function.SqlFunction; -import com.facebook.presto.spi.function.SqlInvokedFunction; -import com.facebook.presto.spi.function.SqlInvokedScalarFunctionImplementation; -import com.google.common.base.Supplier; -import com.google.common.base.Suppliers; -import com.google.common.cache.CacheBuilder; -import com.google.common.cache.CacheLoader; -import com.google.common.cache.LoadingCache; -import com.google.common.util.concurrent.UncheckedExecutionException; import java.util.Collection; import java.util.List; -import java.util.Optional; -import java.util.concurrent.CompletableFuture; import static com.facebook.presto.metadata.BuiltInFunctionKind.PLUGIN; import static com.facebook.presto.spi.function.FunctionImplementationType.SQL; -import static com.facebook.presto.spi.function.FunctionKind.SCALAR; import static com.google.common.base.Preconditions.checkArgument; -import static com.google.common.base.Throwables.throwIfInstanceOf; -import static com.google.common.collect.ImmutableList.toImmutableList; -import static java.util.Collections.emptyList; -import static java.util.Objects.requireNonNull; -import static java.util.concurrent.TimeUnit.HOURS; public class BuiltInPluginFunctionNamespaceManager - implements FunctionNamespaceManager + extends BuiltInSpecialFunctionNamespaceManager { - private volatile FunctionMap functions = new FunctionMap(); - private final FunctionAndTypeManager functionAndTypeManager; - private final Supplier cachedFunctions = - Suppliers.memoize(this::checkForNamingConflicts); - private final LoadingCache specializedFunctionKeyCache; - private final LoadingCache specializedScalarCache; - public BuiltInPluginFunctionNamespaceManager(FunctionAndTypeManager functionAndTypeManager) { - this.functionAndTypeManager = requireNonNull(functionAndTypeManager, "functionAndTypeManager is null"); - specializedFunctionKeyCache = CacheBuilder.newBuilder() - .maximumSize(1000) - .expireAfterWrite(1, HOURS) - .build(CacheLoader.from(this::doGetSpecializedFunctionKey)); - specializedScalarCache = CacheBuilder.newBuilder() - .maximumSize(1000) - .expireAfterWrite(1, HOURS) - .build(CacheLoader.from(key -> { - checkArgument( - key.getFunction() instanceof SqlInvokedFunction, - "Unsupported scalar function class: %s", - key.getFunction().getClass()); - return new SqlInvokedScalarFunctionImplementation(((SqlInvokedFunction) key.getFunction()).getBody()); - })); - } - - public synchronized void registerPluginFunctions(List functions) - { - checkForNamingConflicts(functions); - this.functions = new FunctionMap(this.functions, functions); - } - - @Override - public FunctionHandle getFunctionHandle(Optional transactionHandle, Signature signature) - { - return new BuiltInFunctionHandle(signature, PLUGIN); - } - - @Override - public Collection getFunctions(Optional transactionHandle, QualifiedObjectName functionName) - { - if (functions.list().isEmpty() || - (!functionName.getCatalogSchemaName().equals(functionAndTypeManager.getDefaultNamespace()))) { - return emptyList(); - } - return cachedFunctions.get().get(functionName); - } - - /** - * likePattern / escape is not used for optimization, returning all functions. - */ - @Override - public Collection listFunctions(Optional likePattern, Optional escape) - { - return cachedFunctions.get().list(); - } - - public FunctionMetadata getFunctionMetadata(FunctionHandle functionHandle) - { - checkArgument(functionHandle instanceof BuiltInFunctionHandle, "Expect BuiltInFunctionHandle"); - Signature signature = ((BuiltInFunctionHandle) functionHandle).getSignature(); - SpecializedFunctionKey functionKey; - try { - functionKey = specializedFunctionKeyCache.getUnchecked(signature); - } - catch (UncheckedExecutionException e) { - throwIfInstanceOf(e.getCause(), PrestoException.class); - throw e; - } - SqlFunction function = functionKey.getFunction(); - checkArgument(function instanceof SqlInvokedFunction, "BuiltInPluginFunctionNamespaceManager only support SqlInvokedFunctions"); - SqlInvokedFunction sqlFunction = (SqlInvokedFunction) function; - List argumentNames = sqlFunction.getParameters().stream().map(Parameter::getName).collect(toImmutableList()); - return new FunctionMetadata( - signature.getName(), - signature.getArgumentTypes(), - argumentNames, - signature.getReturnType(), - signature.getKind(), - sqlFunction.getRoutineCharacteristics().getLanguage(), - SQL, - function.isDeterministic(), - function.isCalledOnNullInput(), - sqlFunction.getVersion(), - sqlFunction.getComplexTypeFunctionDescriptor()); - } - - public ScalarFunctionImplementation getScalarFunctionImplementation(FunctionHandle functionHandle) - { - checkArgument(functionHandle instanceof BuiltInFunctionHandle, "Expect BuiltInFunctionHandle"); - return getScalarFunctionImplementation(((BuiltInFunctionHandle) functionHandle).getSignature()); - } - - @Override - public void setBlockEncodingSerde(BlockEncodingSerde blockEncodingSerde) - { - throw new UnsupportedOperationException("BuiltInPluginFunctionNamespaceManager does not support setting block encoding"); - } - - @Override - public FunctionNamespaceTransactionHandle beginTransaction() - { - throw new UnsupportedOperationException("BuiltInPluginFunctionNamespaceManager does not support setting block encoding"); - } - - @Override - public void commit(FunctionNamespaceTransactionHandle transactionHandle) - { - throw new UnsupportedOperationException("BuiltInPluginFunctionNamespaceManager does not support setting block encoding"); - } - - @Override - public void abort(FunctionNamespaceTransactionHandle transactionHandle) - { - throw new UnsupportedOperationException("BuiltInPluginFunctionNamespaceManager does not support setting block encoding"); - } - - @Override - public void createFunction(SqlInvokedFunction function, boolean replace) - { - throw new UnsupportedOperationException("BuiltInPluginFunctionNamespaceManager does not support setting block encoding"); + super(functionAndTypeManager); } - @Override - public void dropFunction(QualifiedObjectName functionName, Optional parameterTypes, boolean exists) - { - throw new UnsupportedOperationException("BuiltInPluginFunctionNamespaceManager does not support drop function"); - } - - @Override - public void alterFunction(QualifiedObjectName functionName, Optional parameterTypes, AlterRoutineCharacteristics alterRoutineCharacteristics) - { - throw new UnsupportedOperationException("BuiltInPluginFunctionNamespaceManager does not alter function"); - } - - @Override - public void addUserDefinedType(UserDefinedType userDefinedType) + public void triggerConflictCheckWithBuiltInFunctions() { - throw new UnsupportedOperationException("BuiltInPluginFunctionNamespaceManager does not support adding user defined types"); + checkForNamingConflicts(this.getFunctionsFromDefaultNamespace()); } @Override - public Optional getUserDefinedType(QualifiedObjectName typeName) + public synchronized void registerBuiltInSpecialFunctions(List functions) { - throw new UnsupportedOperationException("BuiltInPluginFunctionNamespaceManager does not support getting user defined types"); + checkForNamingConflicts(functions); + this.functions = new FunctionMap(this.functions, functions); } @Override - public CompletableFuture executeFunction(String source, FunctionHandle functionHandle, Page input, List channels, TypeManager typeManager) - { - throw new UnsupportedOperationException("BuiltInPluginFunctionNamespaceManager does not execute function"); - } - - private ScalarFunctionImplementation getScalarFunctionImplementation(Signature signature) - { - checkArgument(signature.getKind() == SCALAR, "%s is not a scalar function", signature); - checkArgument(signature.getTypeVariableConstraints().isEmpty(), "%s has unbound type parameters", signature); - - try { - return specializedScalarCache.getUnchecked(getSpecializedFunctionKey(signature)); - } - catch (UncheckedExecutionException e) { - throwIfInstanceOf(e.getCause(), PrestoException.class); - throw e; - } - } - - private synchronized FunctionMap checkForNamingConflicts() - { - Optional> functionNamespaceManager = - functionAndTypeManager.getServingFunctionNamespaceManager(functionAndTypeManager.getDefaultNamespace()); - checkArgument(functionNamespaceManager.isPresent(), "Cannot find function namespace for catalog '%s'", functionAndTypeManager.getDefaultNamespace().getCatalogName()); - checkForNamingConflicts(functionNamespaceManager.get().listFunctions(Optional.empty(), Optional.empty())); - return functions; - } - - private synchronized void checkForNamingConflicts(Collection functions) + protected synchronized void checkForNamingConflicts(Collection functions) { for (SqlFunction function : functions) { for (SqlFunction existingFunction : this.functions.list()) { @@ -241,19 +53,15 @@ private synchronized void checkForNamingConflicts(Collection +{ + protected volatile FunctionMap functions = new FunctionMap(); + private final FunctionAndTypeManager functionAndTypeManager; + private final LoadingCache specializedFunctionKeyCache; + private final LoadingCache specializedScalarCache; + + public BuiltInSpecialFunctionNamespaceManager(FunctionAndTypeManager functionAndTypeManager) + { + this.functionAndTypeManager = requireNonNull(functionAndTypeManager, "functionAndTypeManager is null"); + specializedFunctionKeyCache = CacheBuilder.newBuilder() + .maximumSize(1000) + .expireAfterWrite(1, HOURS) + .build(CacheLoader.from(this::doGetSpecializedFunctionKey)); + specializedScalarCache = CacheBuilder.newBuilder() + .maximumSize(1000) + .expireAfterWrite(1, HOURS) + .build(CacheLoader.from(key -> { + checkArgument( + key.getFunction() instanceof SqlInvokedFunction, + "Unsupported scalar function class: %s", + key.getFunction().getClass()); + return new SqlInvokedScalarFunctionImplementation(((SqlInvokedFunction) key.getFunction()).getBody()); + })); + } + + @Override + public Collection getFunctions(Optional transactionHandle, QualifiedObjectName functionName) + { + return functions.get(functionName); + } + + /** + * likePattern / escape is not used for optimization, returning all functions. + */ + @Override + public Collection listFunctions(Optional likePattern, Optional escape) + { + return functions.list(); + } + + @Override + public FunctionHandle getFunctionHandle(Optional transactionHandle, Signature signature) + { + return new BuiltInFunctionHandle(signature, getBuiltInFunctionKind()); + } + + public FunctionMetadata getFunctionMetadata(FunctionHandle functionHandle) + { + checkArgument(functionHandle instanceof BuiltInFunctionHandle, "Expect BuiltInFunctionHandle"); + Signature signature = ((BuiltInFunctionHandle) functionHandle).getSignature(); + SpecializedFunctionKey functionKey; + try { + functionKey = specializedFunctionKeyCache.getUnchecked(signature); + } + catch (UncheckedExecutionException e) { + throwIfInstanceOf(e.getCause(), PrestoException.class); + throw e; + } + SqlFunction function = functionKey.getFunction(); + checkArgument(function instanceof SqlInvokedFunction, "BuiltInPluginFunctionNamespaceManager only support SqlInvokedFunctions"); + SqlInvokedFunction sqlFunction = (SqlInvokedFunction) function; + List argumentNames = sqlFunction.getParameters().stream().map(Parameter::getName).collect(toImmutableList()); + return new FunctionMetadata( + signature.getName(), + signature.getArgumentTypes(), + argumentNames, + signature.getReturnType(), + signature.getKind(), + sqlFunction.getRoutineCharacteristics().getLanguage(), + getDefaultFunctionMetadataImplementationType(), + function.isDeterministic(), + function.isCalledOnNullInput(), + sqlFunction.getVersion(), + sqlFunction.getComplexTypeFunctionDescriptor()); + } + + @Override + public void setBlockEncodingSerde(BlockEncodingSerde blockEncodingSerde) + { + throw new UnsupportedOperationException("BuiltInPluginFunctionNamespaceManager does not support setting block encoding"); + } + + @Override + public FunctionNamespaceTransactionHandle beginTransaction() + { + throw new UnsupportedOperationException("BuiltInPluginFunctionNamespaceManager does not support setting block encoding"); + } + + @Override + public void commit(FunctionNamespaceTransactionHandle transactionHandle) + { + throw new UnsupportedOperationException("BuiltInPluginFunctionNamespaceManager does not support setting block encoding"); + } + + @Override + public void abort(FunctionNamespaceTransactionHandle transactionHandle) + { + throw new UnsupportedOperationException("BuiltInPluginFunctionNamespaceManager does not support setting block encoding"); + } + + @Override + public void createFunction(SqlInvokedFunction function, boolean replace) + { + throw new UnsupportedOperationException("BuiltInPluginFunctionNamespaceManager does not support setting block encoding"); + } + + @Override + public void dropFunction(QualifiedObjectName functionName, Optional parameterTypes, boolean exists) + { + throw new UnsupportedOperationException("BuiltInPluginFunctionNamespaceManager does not support drop function"); + } + + @Override + public void alterFunction(QualifiedObjectName functionName, Optional parameterTypes, AlterRoutineCharacteristics alterRoutineCharacteristics) + { + throw new UnsupportedOperationException("BuiltInPluginFunctionNamespaceManager does not alter function"); + } + + @Override + public void addUserDefinedType(UserDefinedType userDefinedType) + { + throw new UnsupportedOperationException("BuiltInPluginFunctionNamespaceManager does not support adding user defined types"); + } + + @Override + public Optional getUserDefinedType(QualifiedObjectName typeName) + { + throw new UnsupportedOperationException("BuiltInPluginFunctionNamespaceManager does not support getting user defined types"); + } + + @Override + public CompletableFuture executeFunction(String source, FunctionHandle functionHandle, Page input, List channels, TypeManager typeManager) + { + throw new UnsupportedOperationException("BuiltInPluginFunctionNamespaceManager does not execute function"); + } + + protected abstract void checkForNamingConflicts(Collection functions); + + protected abstract BuiltInFunctionKind getBuiltInFunctionKind(); + + protected abstract FunctionImplementationType getDefaultFunctionMetadataImplementationType(); + + public abstract void registerBuiltInSpecialFunctions(List functions); + + @Override + public ScalarFunctionImplementation getScalarFunctionImplementation(FunctionHandle functionHandle) + { + checkArgument(functionHandle instanceof BuiltInFunctionHandle, "Expect BuiltInFunctionHandle"); + return getScalarFunctionImplementation(((BuiltInFunctionHandle) functionHandle).getSignature()); + } + + protected Collection getFunctionsFromDefaultNamespace() + { + Optional> functionNamespaceManager = + functionAndTypeManager.getServingFunctionNamespaceManager(functionAndTypeManager.getDefaultNamespace()); + checkArgument(functionNamespaceManager.isPresent(), "Cannot find function namespace for catalog '%s'", functionAndTypeManager.getDefaultNamespace().getCatalogName()); + return functionNamespaceManager.get().listFunctions(Optional.empty(), Optional.empty()); + } + + private synchronized FunctionMap createFunctionMap() + { + return functions; + } + + private ScalarFunctionImplementation getScalarFunctionImplementation(Signature signature) + { + checkArgument(signature.getKind() == SCALAR, "%s is not a scalar function", signature); + checkArgument(signature.getTypeVariableConstraints().isEmpty(), "%s has unbound type parameters", signature); + + try { + return specializedScalarCache.getUnchecked(getSpecializedFunctionKey(signature)); + } + catch (UncheckedExecutionException e) { + throwIfInstanceOf(e.getCause(), PrestoException.class); + throw e; + } + } + + private SpecializedFunctionKey doGetSpecializedFunctionKey(Signature signature) + { + return functionAndTypeManager.getSpecializedFunctionKey(signature, getFunctions(Optional.empty(), signature.getName())); + } + + private SpecializedFunctionKey getSpecializedFunctionKey(Signature signature) + { + try { + return specializedFunctionKeyCache.getUnchecked(signature); + } + catch (UncheckedExecutionException e) { + throwIfInstanceOf(e.getCause(), PrestoException.class); + throw e; + } + } +} diff --git a/presto-main-base/src/main/java/com/facebook/presto/metadata/BuiltInTypeAndFunctionNamespaceManager.java b/presto-main-base/src/main/java/com/facebook/presto/metadata/BuiltInTypeAndFunctionNamespaceManager.java index f7565a43b47d2..4984cd8fe065d 100644 --- a/presto-main-base/src/main/java/com/facebook/presto/metadata/BuiltInTypeAndFunctionNamespaceManager.java +++ b/presto-main-base/src/main/java/com/facebook/presto/metadata/BuiltInTypeAndFunctionNamespaceManager.java @@ -13,7 +13,6 @@ */ package com.facebook.presto.metadata; -import com.facebook.presto.UnknownTypeException; import com.facebook.presto.common.CatalogSchemaName; import com.facebook.presto.common.Page; import com.facebook.presto.common.QualifiedObjectName; @@ -68,6 +67,7 @@ import com.facebook.presto.operator.aggregation.DoubleCovarianceAggregation; import com.facebook.presto.operator.aggregation.DoubleHistogramAggregation; import com.facebook.presto.operator.aggregation.DoubleRegressionAggregation; +import com.facebook.presto.operator.aggregation.DoubleRegressionExtendedAggregation; import com.facebook.presto.operator.aggregation.DoubleSumAggregation; import com.facebook.presto.operator.aggregation.EntropyAggregation; import com.facebook.presto.operator.aggregation.GeometricMeanAggregations; @@ -85,6 +85,7 @@ import com.facebook.presto.operator.aggregation.RealGeometricMeanAggregations; import com.facebook.presto.operator.aggregation.RealHistogramAggregation; import com.facebook.presto.operator.aggregation.RealRegressionAggregation; +import com.facebook.presto.operator.aggregation.RealRegressionExtendedAggregation; import com.facebook.presto.operator.aggregation.RealSumAggregation; import com.facebook.presto.operator.aggregation.ReduceAggregationFunction; import com.facebook.presto.operator.aggregation.SumDataSizeForStats; @@ -109,6 +110,7 @@ import com.facebook.presto.operator.aggregation.sketch.kll.KllSketchAggregationFunction; import com.facebook.presto.operator.aggregation.sketch.kll.KllSketchWithKAggregationFunction; import com.facebook.presto.operator.aggregation.sketch.theta.ThetaSketchAggregationFunction; +import com.facebook.presto.operator.scalar.AbstractArraySortByKeyFunction; import com.facebook.presto.operator.scalar.ArrayAllMatchFunction; import com.facebook.presto.operator.scalar.ArrayAnyMatchFunction; import com.facebook.presto.operator.scalar.ArrayCardinalityFunction; @@ -204,11 +206,6 @@ import com.facebook.presto.operator.scalar.WilsonInterval; import com.facebook.presto.operator.scalar.WordStemFunction; import com.facebook.presto.operator.scalar.queryplan.JsonPrestoQueryPlanFunctions; -import com.facebook.presto.operator.scalar.sql.ArraySqlFunctions; -import com.facebook.presto.operator.scalar.sql.MapNormalizeFunction; -import com.facebook.presto.operator.scalar.sql.MapSqlFunctions; -import com.facebook.presto.operator.scalar.sql.SimpleSamplingPercent; -import com.facebook.presto.operator.scalar.sql.StringSqlFunctions; import com.facebook.presto.operator.window.CumulativeDistributionFunction; import com.facebook.presto.operator.window.DenseRankFunction; import com.facebook.presto.operator.window.FirstValueFunction; @@ -236,6 +233,7 @@ import com.facebook.presto.spi.function.SqlFunctionVisibility; import com.facebook.presto.spi.function.SqlInvokedFunction; import com.facebook.presto.spi.function.SqlInvokedScalarFunctionImplementation; +import com.facebook.presto.spi.type.UnknownTypeException; import com.facebook.presto.sql.analyzer.FunctionsConfig; import com.facebook.presto.type.BigintOperators; import com.facebook.presto.type.BooleanOperators; @@ -749,7 +747,9 @@ private List getBuiltInFunctions(FunctionsConfig function .aggregates(DoubleCovarianceAggregation.class) .aggregates(RealCovarianceAggregation.class) .aggregates(DoubleRegressionAggregation.class) + .aggregates(DoubleRegressionExtendedAggregation.class) .aggregates(RealRegressionAggregation.class) + .aggregates(RealRegressionExtendedAggregation.class) .aggregates(DoubleCorrelationAggregation.class) .aggregates(RealCorrelationAggregation.class) .aggregates(BitwiseOrAggregation.class) @@ -883,6 +883,8 @@ private List getBuiltInFunctions(FunctionsConfig function .scalar(ArrayGreaterThanOrEqualOperator.class) .scalar(ArrayElementAtFunction.class) .scalar(ArraySortFunction.class) + .function(AbstractArraySortByKeyFunction.ArraySortByKeyFunction.ARRAY_SORT_BY_KEY_FUNCTION) + .function(AbstractArraySortByKeyFunction.ArraySortDescByKeyFunction.ARRAY_SORT_DESC_BY_KEY_FUNCTION) .scalar(MapSubsetFunction.class) .scalar(ArraySortComparatorFunction.class) .scalar(ArrayShuffleFunction.class) @@ -994,12 +996,6 @@ private List getBuiltInFunctions(FunctionsConfig function .aggregate(ThetaSketchAggregationFunction.class) .scalars(ThetaSketchFunctions.class) .function(MergeTDigestFunction.MERGE) - .sqlInvokedScalar(MapNormalizeFunction.class) - .sqlInvokedScalars(ArraySqlFunctions.class) - .sqlInvokedScalars(ArrayIntersectFunction.class) - .sqlInvokedScalars(MapSqlFunctions.class) - .sqlInvokedScalars(SimpleSamplingPercent.class) - .sqlInvokedScalars(StringSqlFunctions.class) .scalar(DynamicFilterPlaceholderFunction.class) .scalars(EnumCasts.class) .scalars(LongEnumOperators.class) @@ -1285,6 +1281,18 @@ public Type getType(TypeSignature typeSignature) } } + @Override + public boolean hasType(TypeSignature typeSignature) + { + try { + getType(typeSignature); + return true; + } + catch (UnknownTypeException e) { + return false; + } + } + @Override public Type getParameterizedType(String baseTypeName, List typeParameters) { diff --git a/presto-main-base/src/main/java/com/facebook/presto/metadata/BuiltInWorkerFunctionNamespaceManager.java b/presto-main-base/src/main/java/com/facebook/presto/metadata/BuiltInWorkerFunctionNamespaceManager.java new file mode 100644 index 0000000000000..41ef943001cb3 --- /dev/null +++ b/presto-main-base/src/main/java/com/facebook/presto/metadata/BuiltInWorkerFunctionNamespaceManager.java @@ -0,0 +1,59 @@ +/* + * Licensed 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 com.facebook.presto.metadata; + +import com.facebook.presto.spi.function.FunctionImplementationType; +import com.facebook.presto.spi.function.SqlFunction; + +import java.util.Collection; +import java.util.List; + +import static com.facebook.presto.metadata.BuiltInFunctionKind.WORKER; +import static com.facebook.presto.spi.function.FunctionImplementationType.CPP; + +public class BuiltInWorkerFunctionNamespaceManager + extends BuiltInSpecialFunctionNamespaceManager +{ + public BuiltInWorkerFunctionNamespaceManager(FunctionAndTypeManager functionAndTypeManager) + { + super(functionAndTypeManager); + } + + @Override + public synchronized void registerBuiltInSpecialFunctions(List functions) + { + // only register functions once + if (!this.functions.list().isEmpty()) { + return; + } + this.functions = new FunctionMap(this.functions, functions); + } + + @Override + protected synchronized void checkForNamingConflicts(Collection functions) + { + } + + @Override + protected BuiltInFunctionKind getBuiltInFunctionKind() + { + return WORKER; + } + + @Override + protected FunctionImplementationType getDefaultFunctionMetadataImplementationType() + { + return CPP; + } +} diff --git a/presto-main-base/src/main/java/com/facebook/presto/metadata/Catalog.java b/presto-main-base/src/main/java/com/facebook/presto/metadata/Catalog.java index 4ee7ed05f963b..5cbdb9bfb1308 100644 --- a/presto-main-base/src/main/java/com/facebook/presto/metadata/Catalog.java +++ b/presto-main-base/src/main/java/com/facebook/presto/metadata/Catalog.java @@ -32,6 +32,8 @@ public class Catalog private final ConnectorId systemTablesId; private final Connector systemTables; + private final CatalogContext catalogContext; + public Catalog( String catalogName, ConnectorId connectorId, @@ -40,6 +42,26 @@ public Catalog( Connector informationSchema, ConnectorId systemTablesId, Connector systemTables) + { + this(catalogName, + connectorId, + connector, + informationSchemaId, + informationSchema, + systemTablesId, + systemTables, + catalogName); + } + + public Catalog( + String catalogName, + ConnectorId connectorId, + Connector connector, + ConnectorId informationSchemaId, + Connector informationSchema, + ConnectorId systemTablesId, + Connector systemTables, + String connectorName) { this.catalogName = checkCatalogName(catalogName); this.connectorId = requireNonNull(connectorId, "connectorId is null"); @@ -48,8 +70,9 @@ public Catalog( this.informationSchema = requireNonNull(informationSchema, "informationSchema is null"); this.systemTablesId = requireNonNull(systemTablesId, "systemTablesId is null"); this.systemTables = requireNonNull(systemTables, "systemTables is null"); + requireNonNull(connectorName, "connectorName is null"); + this.catalogContext = new CatalogContext(catalogName, connectorName); } - public String getCatalogName() { return catalogName; @@ -60,6 +83,11 @@ public ConnectorId getConnectorId() return connectorId; } + public CatalogContext getCatalogContext() + { + return catalogContext; + } + public ConnectorId getInformationSchemaId() { return informationSchemaId; @@ -92,4 +120,26 @@ public String toString() .add("connectorId", connectorId) .toString(); } + + public class CatalogContext + { + private final String catalogName; + private final String connectorName; + + public CatalogContext(String catalogName, String connectorName) + { + this.catalogName = catalogName; + this.connectorName = connectorName; + } + + public String getCatalogName() + { + return catalogName; + } + + public String getConnectorName() + { + return connectorName; + } + } } diff --git a/presto-main-base/src/main/java/com/facebook/presto/metadata/DelegatingMetadataManager.java b/presto-main-base/src/main/java/com/facebook/presto/metadata/DelegatingMetadataManager.java index 7302236671319..818e342cd8887 100644 --- a/presto-main-base/src/main/java/com/facebook/presto/metadata/DelegatingMetadataManager.java +++ b/presto-main-base/src/main/java/com/facebook/presto/metadata/DelegatingMetadataManager.java @@ -34,6 +34,7 @@ import com.facebook.presto.spi.connector.ConnectorCapabilities; import com.facebook.presto.spi.connector.ConnectorOutputMetadata; import com.facebook.presto.spi.connector.ConnectorTableVersion; +import com.facebook.presto.spi.connector.TableFunctionApplicationResult; import com.facebook.presto.spi.constraints.TableConstraint; import com.facebook.presto.spi.function.SqlFunction; import com.facebook.presto.spi.plan.PartitioningHandle; @@ -657,4 +658,10 @@ public String normalizeIdentifier(Session session, String catalogName, String id { return delegate.normalizeIdentifier(session, catalogName, identifier); } + + @Override + public Optional> applyTableFunction(Session session, TableFunctionHandle handle) + { + return delegate.applyTableFunction(session, handle); + } } diff --git a/presto-main-base/src/main/java/com/facebook/presto/metadata/FunctionAndTypeManager.java b/presto-main-base/src/main/java/com/facebook/presto/metadata/FunctionAndTypeManager.java index a409fdddb4afc..9f18bbed9fc00 100644 --- a/presto-main-base/src/main/java/com/facebook/presto/metadata/FunctionAndTypeManager.java +++ b/presto-main-base/src/main/java/com/facebook/presto/metadata/FunctionAndTypeManager.java @@ -13,6 +13,7 @@ */ package com.facebook.presto.metadata; +import com.facebook.airlift.log.Logger; import com.facebook.presto.Session; import com.facebook.presto.common.CatalogSchemaName; import com.facebook.presto.common.Page; @@ -39,6 +40,7 @@ import com.facebook.presto.spi.function.AggregationFunctionImplementation; import com.facebook.presto.spi.function.AlterRoutineCharacteristics; import com.facebook.presto.spi.function.FunctionHandle; +import com.facebook.presto.spi.function.FunctionImplementationType; import com.facebook.presto.spi.function.FunctionMetadata; import com.facebook.presto.spi.function.FunctionMetadataManager; import com.facebook.presto.spi.function.FunctionNamespaceManager; @@ -94,6 +96,7 @@ import static com.facebook.presto.SystemSessionProperties.isListBuiltInFunctionsOnly; import static com.facebook.presto.common.type.TypeSignature.parseTypeSignature; import static com.facebook.presto.metadata.BuiltInFunctionKind.PLUGIN; +import static com.facebook.presto.metadata.BuiltInFunctionKind.WORKER; import static com.facebook.presto.metadata.BuiltInTypeAndFunctionNamespaceManager.JAVA_BUILTIN_NAMESPACE; import static com.facebook.presto.metadata.CastType.toOperatorType; import static com.facebook.presto.metadata.FunctionSignatureMatcher.constructFunctionNotFoundErrorMessage; @@ -132,6 +135,7 @@ public class FunctionAndTypeManager implements FunctionMetadataManager, TypeManager { private static final Pattern DEFAULT_NAMESPACE_PREFIX_PATTERN = Pattern.compile("[a-z]+\\.[a-z]+"); + private static final Logger log = Logger.get(FunctionAndTypeManager.class); private final TransactionManager transactionManager; private final TableFunctionRegistry tableFunctionRegistry; private final BlockEncodingSerde blockEncodingSerde; @@ -147,9 +151,11 @@ public class FunctionAndTypeManager private final LoadingCache functionCache; private final CacheStatsMBean cacheStatsMBean; private final boolean nativeExecution; + private final boolean isBuiltInSidecarFunctionsEnabled; private final CatalogSchemaName defaultNamespace; private final AtomicReference servingTypeManager; private final AtomicReference>> servingTypeManagerParametricTypesSupplier; + private final BuiltInWorkerFunctionNamespaceManager builtInWorkerFunctionNamespaceManager; private final BuiltInPluginFunctionNamespaceManager builtInPluginFunctionNamespaceManager; private final FunctionsConfig functionsConfig; private final Set types; @@ -185,9 +191,11 @@ public FunctionAndTypeManager( this.functionSignatureMatcher = new FunctionSignatureMatcher(this); this.typeCoercer = new TypeCoercer(functionsConfig, this); this.nativeExecution = featuresConfig.isNativeExecutionEnabled(); + this.isBuiltInSidecarFunctionsEnabled = featuresConfig.isBuiltInSidecarFunctionsEnabled(); this.defaultNamespace = configureDefaultNamespace(functionsConfig.getDefaultNamespacePrefix()); this.servingTypeManager = new AtomicReference<>(builtInTypeAndFunctionNamespaceManager); this.servingTypeManagerParametricTypesSupplier = new AtomicReference<>(this::getServingTypeManagerParametricTypes); + this.builtInWorkerFunctionNamespaceManager = new BuiltInWorkerFunctionNamespaceManager(this); this.builtInPluginFunctionNamespaceManager = new BuiltInPluginFunctionNamespaceManager(this); } @@ -278,6 +286,12 @@ public Collection getParametricTypes() return FunctionAndTypeManager.this.getParametricTypes(); } + @Override + public boolean hasType(TypeSignature signature) + { + return FunctionAndTypeManager.this.hasType(signature); + } + @Override public Collection listBuiltInFunctions() { @@ -361,8 +375,12 @@ public FunctionMetadata getFunctionMetadata(FunctionHandle functionHandle) if (isBuiltInPluginFunctionHandle(functionHandle)) { return builtInPluginFunctionNamespaceManager.getFunctionMetadata(functionHandle); } + if (isBuiltInWorkerFunctionHandle(functionHandle)) { + return builtInWorkerFunctionNamespaceManager.getFunctionMetadata(functionHandle); + } Optional> functionNamespaceManager = getServingFunctionNamespaceManager(functionHandle.getCatalogSchemaName()); checkArgument(functionNamespaceManager.isPresent(), "Cannot find function namespace for '%s'", functionHandle.getCatalogSchemaName()); + return functionNamespaceManager.get().getFunctionMetadata(functionHandle); } @@ -395,6 +413,12 @@ public Type getType(TypeSignature signature) return getUserDefinedType(signature); } + @Override + public boolean hasType(TypeSignature signature) + { + return servingTypeManager.get().hasType(signature); + } + @Override public Type getParameterizedType(String baseTypeName, List typeParameters) { @@ -452,14 +476,26 @@ public void addTypeManagerFactory(TypeManagerFactory factory) } } + public TransactionManager getTransactionManager() + { + return transactionManager; + } + public void registerBuiltInFunctions(List functions) { builtInTypeAndFunctionNamespaceManager.registerBuiltInFunctions(functions); } + public void registerWorkerFunctions(List functions) + { + if (isBuiltInSidecarFunctionsEnabled) { + builtInWorkerFunctionNamespaceManager.registerBuiltInSpecialFunctions(functions); + } + } + public void registerPluginFunctions(List functions) { - builtInPluginFunctionNamespaceManager.registerPluginFunctions(functions); + builtInPluginFunctionNamespaceManager.registerBuiltInSpecialFunctions(functions); } public void registerConnectorFunctions(String catalogName, List functions) @@ -490,6 +526,7 @@ public List listFunctions(Session session, Optional likePat defaultNamespace.getCatalogName()).listFunctions(likePattern, escape).stream() .collect(toImmutableList())); functions.addAll(builtInPluginFunctionNamespaceManager.listFunctions(likePattern, escape).stream().collect(toImmutableList())); + functions.addAll(builtInWorkerFunctionNamespaceManager.listFunctions(likePattern, escape).stream().collect(toImmutableList())); } else { functions.addAll(SessionFunctionUtils.listFunctions(session.getSessionFunctions())); @@ -497,6 +534,7 @@ public List listFunctions(Session session, Optional likePat .flatMap(manager -> manager.listFunctions(likePattern, escape).stream()) .collect(toImmutableList())); functions.addAll(builtInPluginFunctionNamespaceManager.listFunctions(likePattern, escape).stream().collect(toImmutableList())); + functions.addAll(builtInWorkerFunctionNamespaceManager.listFunctions(likePattern, escape).stream().collect(toImmutableList())); } return functions.build().stream() @@ -642,6 +680,10 @@ public ScalarFunctionImplementation getScalarFunctionImplementation(FunctionHand if (isBuiltInPluginFunctionHandle(functionHandle)) { return builtInPluginFunctionNamespaceManager.getScalarFunctionImplementation(functionHandle); } + if (isBuiltInWorkerFunctionHandle(functionHandle)) { + return builtInWorkerFunctionNamespaceManager.getScalarFunctionImplementation(functionHandle); + } + Optional> functionNamespaceManager = getServingFunctionNamespaceManager(functionHandle.getCatalogSchemaName()); checkArgument(functionNamespaceManager.isPresent(), "Cannot find function namespace for '%s'", functionHandle.getCatalogSchemaName()); return functionNamespaceManager.get().getScalarFunctionImplementation(functionHandle); @@ -942,7 +984,12 @@ public SpecializedFunctionKey getSpecializedFunctionKey(Signature signature, Col return builtInTypeAndFunctionNamespaceManager.doGetSpecializedFunctionKeyForMagicLiteralFunctions(signature, this); } - public CatalogSchemaName configureDefaultNamespace(String defaultNamespacePrefixString) + public BuiltInPluginFunctionNamespaceManager getBuiltInPluginFunctionNamespaceManager() + { + return builtInPluginFunctionNamespaceManager; + } + + private CatalogSchemaName configureDefaultNamespace(String defaultNamespacePrefixString) { if (!defaultNamespacePrefixString.matches(DEFAULT_NAMESPACE_PREFIX_PATTERN.pattern())) { throw new PrestoException(GENERIC_USER_ERROR, format("Default namespace prefix string should be in the form of 'catalog.schema', found: %s", defaultNamespacePrefixString)); @@ -965,13 +1012,15 @@ private Collection getFunctions( return ImmutableList.builder() .addAll(functionNamespaceManager.getFunctions(transactionHandle, functionName)) .addAll(builtInPluginFunctionNamespaceManager.getFunctions(transactionHandle, functionName)) + .addAll(builtInWorkerFunctionNamespaceManager.getFunctions(transactionHandle, functionName)) .build(); } /** * Gets the function handle of the function if there is a match. We enforce explicit naming for dynamic function namespaces. * All unqualified function names will only be resolved against the built-in default function namespace. We get all the candidates - * from the current default namespace and additionally all the candidates from builtInPluginFunctionNamespaceManager. + * from the current default namespace and additionally all the candidates from builtInPluginFunctionNamespaceManager and + * builtInWorkerFunctionNamespaceManager. * * @throws PrestoException if there are no matches or multiple matches */ @@ -986,20 +1035,43 @@ private FunctionHandle getMatchingFunctionHandle( getMatchingFunction(functionNamespaceManager.getFunctions(transactionHandle, functionName), parameterTypes, coercionAllowed); Optional matchingPluginFunctionSignature = getMatchingFunction(builtInPluginFunctionNamespaceManager.getFunctions(transactionHandle, functionName), parameterTypes, coercionAllowed); + Optional matchingWorkerFunctionSignature = + getMatchingFunction(builtInWorkerFunctionNamespaceManager.getFunctions(transactionHandle, functionName), parameterTypes, coercionAllowed); if (matchingDefaultFunctionSignature.isPresent() && matchingPluginFunctionSignature.isPresent()) { throw new PrestoException(AMBIGUOUS_FUNCTION_CALL, format("Function '%s' has two matching signatures. Please specify parameter types. \n" + "First match : '%s', Second match: '%s'", functionName, matchingDefaultFunctionSignature.get(), matchingPluginFunctionSignature.get())); } - if (matchingDefaultFunctionSignature.isPresent()) { - return functionNamespaceManager.getFunctionHandle(transactionHandle, matchingDefaultFunctionSignature.get()); + if (matchingDefaultFunctionSignature.isPresent() && matchingWorkerFunctionSignature.isPresent()) { + FunctionHandle defaultFunctionHandle = functionNamespaceManager.getFunctionHandle(transactionHandle, matchingDefaultFunctionSignature.get()); + FunctionHandle workerFunctionHandle = builtInWorkerFunctionNamespaceManager.getFunctionHandle(transactionHandle, matchingWorkerFunctionSignature.get()); + + if (functionNamespaceManager.getFunctionMetadata(defaultFunctionHandle).getImplementationType().equals(FunctionImplementationType.JAVA)) { + return defaultFunctionHandle; + } + if (functionNamespaceManager.getFunctionMetadata(defaultFunctionHandle).getImplementationType().equals(FunctionImplementationType.SQL)) { + return workerFunctionHandle; + } + } + + if (matchingPluginFunctionSignature.isPresent() && matchingWorkerFunctionSignature.isPresent()) { + // built in plugin function namespace manager always has SQL as implementation type + return builtInWorkerFunctionNamespaceManager.getFunctionHandle(transactionHandle, matchingWorkerFunctionSignature.get()); + } + + if (matchingWorkerFunctionSignature.isPresent()) { + return builtInWorkerFunctionNamespaceManager.getFunctionHandle(transactionHandle, matchingWorkerFunctionSignature.get()); } if (matchingPluginFunctionSignature.isPresent()) { return builtInPluginFunctionNamespaceManager.getFunctionHandle(transactionHandle, matchingPluginFunctionSignature.get()); } + if (matchingDefaultFunctionSignature.isPresent()) { + return functionNamespaceManager.getFunctionHandle(transactionHandle, matchingDefaultFunctionSignature.get()); + } + throw new PrestoException(FUNCTION_NOT_FOUND, constructFunctionNotFoundErrorMessage(functionName, parameterTypes, getFunctions(functionName, transactionHandle, functionNamespaceManager))); } @@ -1017,6 +1089,11 @@ private boolean isBuiltInPluginFunctionHandle(FunctionHandle functionHandle) return (functionHandle instanceof BuiltInFunctionHandle) && ((BuiltInFunctionHandle) functionHandle).getBuiltInFunctionKind().equals(PLUGIN); } + private boolean isBuiltInWorkerFunctionHandle(FunctionHandle functionHandle) + { + return (functionHandle instanceof BuiltInFunctionHandle) && ((BuiltInFunctionHandle) functionHandle).getBuiltInFunctionKind().equals(WORKER); + } + private static class FunctionResolutionCacheKey { private final QualifiedObjectName functionName; diff --git a/presto-main-base/src/main/java/com/facebook/presto/metadata/FunctionListBuilder.java b/presto-main-base/src/main/java/com/facebook/presto/metadata/FunctionListBuilder.java index f3dab976902a9..799dd5f8074fe 100644 --- a/presto-main-base/src/main/java/com/facebook/presto/metadata/FunctionListBuilder.java +++ b/presto-main-base/src/main/java/com/facebook/presto/metadata/FunctionListBuilder.java @@ -15,7 +15,6 @@ import com.facebook.presto.operator.scalar.annotations.CodegenScalarFromAnnotationsParser; import com.facebook.presto.operator.scalar.annotations.ScalarFromAnnotationsParser; -import com.facebook.presto.operator.scalar.annotations.SqlInvokedScalarFromAnnotationsParser; import com.facebook.presto.operator.window.WindowAnnotationsParser; import com.facebook.presto.spi.function.SqlFunction; import com.facebook.presto.spi.function.WindowFunction; @@ -24,7 +23,6 @@ import java.util.ArrayList; import java.util.List; -import static com.facebook.presto.metadata.BuiltInTypeAndFunctionNamespaceManager.JAVA_BUILTIN_NAMESPACE; import static java.util.Objects.requireNonNull; public class FunctionListBuilder @@ -61,18 +59,6 @@ public FunctionListBuilder scalars(Class clazz) return this; } - public FunctionListBuilder sqlInvokedScalar(Class clazz) - { - functions.addAll(SqlInvokedScalarFromAnnotationsParser.parseFunctionDefinition(clazz, JAVA_BUILTIN_NAMESPACE)); - return this; - } - - public FunctionListBuilder sqlInvokedScalars(Class clazz) - { - functions.addAll(SqlInvokedScalarFromAnnotationsParser.parseFunctionDefinitions(clazz, JAVA_BUILTIN_NAMESPACE)); - return this; - } - public FunctionListBuilder codegenScalars(Class clazz) { functions.addAll(CodegenScalarFromAnnotationsParser.parseFunctionDefinitions(clazz)); diff --git a/presto-main-base/src/main/java/com/facebook/presto/metadata/FunctionSignatureMatcher.java b/presto-main-base/src/main/java/com/facebook/presto/metadata/FunctionSignatureMatcher.java index 39cb60c65460a..46d34dc13d61b 100644 --- a/presto-main-base/src/main/java/com/facebook/presto/metadata/FunctionSignatureMatcher.java +++ b/presto-main-base/src/main/java/com/facebook/presto/metadata/FunctionSignatureMatcher.java @@ -71,7 +71,7 @@ public Optional match(Collection candidates, L .filter(function -> !function.getSignature().getTypeVariableConstraints().isEmpty()) .collect(Collectors.toList()); - match = matchFunctionExact(genericCandidates, parameterTypes); + match = matchFunctionGeneric(genericCandidates, parameterTypes); if (match.isPresent()) { return match; } @@ -91,6 +91,28 @@ private Optional matchFunctionExact(List candidates, Lis return matchFunction(candidates, actualParameters, false); } + private Optional matchFunctionGeneric(List candidates, List actualParameters) + { + List applicableFunctions = identifyApplicableFunctions(candidates, actualParameters, false); + if (applicableFunctions.isEmpty()) { + return Optional.empty(); + } + + if (applicableFunctions.size() == 1) { + return Optional.of(getOnlyElement(applicableFunctions).getBoundSignature()); + } + + List deduplicatedSignatures = applicableFunctions.stream() + .map(applicableFunction -> applicableFunction.boundSignature) + .distinct() + .collect(toImmutableList()); + if (deduplicatedSignatures.size() == 1) { + return Optional.of(getOnlyElement(deduplicatedSignatures)); + } + + throw new PrestoException(AMBIGUOUS_FUNCTION_CALL, getErrorMessage(applicableFunctions)); + } + private Optional matchFunctionWithCoercion(Collection candidates, List actualParameters) { return matchFunction(candidates, actualParameters, true); @@ -112,6 +134,11 @@ private Optional matchFunction(Collection cand return Optional.of(getOnlyElement(applicableFunctions).getBoundSignature()); } + throw new PrestoException(AMBIGUOUS_FUNCTION_CALL, getErrorMessage(applicableFunctions)); + } + + private String getErrorMessage(List applicableFunctions) + { StringBuilder errorMessageBuilder = new StringBuilder(); errorMessageBuilder.append("Could not choose a best candidate operator. Explicit type casts must be added.\n"); errorMessageBuilder.append("Candidates are:\n"); @@ -120,7 +147,8 @@ private Optional matchFunction(Collection cand errorMessageBuilder.append(function.getBoundSignature().toString()); errorMessageBuilder.append("\n"); } - throw new PrestoException(AMBIGUOUS_FUNCTION_CALL, errorMessageBuilder.toString()); + + return errorMessageBuilder.toString(); } private List identifyApplicableFunctions(Collection candidates, List actualParameters, boolean allowCoercion) diff --git a/presto-main-base/src/main/java/com/facebook/presto/metadata/InMemoryNodeManager.java b/presto-main-base/src/main/java/com/facebook/presto/metadata/InMemoryNodeManager.java index 69ed24cb1c5bd..f0a3ca28578fa 100644 --- a/presto-main-base/src/main/java/com/facebook/presto/metadata/InMemoryNodeManager.java +++ b/presto-main-base/src/main/java/com/facebook/presto/metadata/InMemoryNodeManager.java @@ -15,6 +15,7 @@ import com.facebook.presto.client.NodeVersion; import com.facebook.presto.spi.ConnectorId; +import com.facebook.presto.spi.NodeLoadMetrics; import com.facebook.presto.spi.NodeState; import com.google.common.collect.HashMultimap; import com.google.common.collect.ImmutableList; @@ -27,6 +28,7 @@ import java.net.URI; import java.util.ArrayList; import java.util.List; +import java.util.Optional; import java.util.Set; import java.util.function.Consumer; import java.util.stream.Stream; @@ -179,4 +181,10 @@ public synchronized void removeNodeChangeListener(Consumer listener) { listeners.remove(requireNonNull(listener, "listener is null")); } + + @Override + public Optional getNodeLoadMetrics(String nodeIdentifier) + { + return Optional.empty(); + } } diff --git a/presto-main-base/src/main/java/com/facebook/presto/metadata/InternalNodeManager.java b/presto-main-base/src/main/java/com/facebook/presto/metadata/InternalNodeManager.java index 81fa60c8165d7..c080dc40fd42d 100644 --- a/presto-main-base/src/main/java/com/facebook/presto/metadata/InternalNodeManager.java +++ b/presto-main-base/src/main/java/com/facebook/presto/metadata/InternalNodeManager.java @@ -14,8 +14,10 @@ package com.facebook.presto.metadata; import com.facebook.presto.spi.ConnectorId; +import com.facebook.presto.spi.NodeLoadMetrics; import com.facebook.presto.spi.NodeState; +import java.util.Optional; import java.util.Set; import java.util.function.Consumer; @@ -46,4 +48,6 @@ public interface InternalNodeManager void addNodeChangeListener(Consumer listener); void removeNodeChangeListener(Consumer listener); + + Optional getNodeLoadMetrics(String nodeIdentifier); } diff --git a/presto-main-base/src/main/java/com/facebook/presto/metadata/Metadata.java b/presto-main-base/src/main/java/com/facebook/presto/metadata/Metadata.java index f7283b436c65f..f53d1b7022746 100644 --- a/presto-main-base/src/main/java/com/facebook/presto/metadata/Metadata.java +++ b/presto-main-base/src/main/java/com/facebook/presto/metadata/Metadata.java @@ -20,6 +20,7 @@ import com.facebook.presto.common.predicate.TupleDomain; import com.facebook.presto.common.type.Type; import com.facebook.presto.common.type.TypeSignature; +import com.facebook.presto.metadata.Catalog.CatalogContext; import com.facebook.presto.spi.ColumnHandle; import com.facebook.presto.spi.ColumnMetadata; import com.facebook.presto.spi.ConnectorId; @@ -40,6 +41,7 @@ import com.facebook.presto.spi.connector.ConnectorOutputMetadata; import com.facebook.presto.spi.connector.ConnectorPartitioningHandle; import com.facebook.presto.spi.connector.ConnectorTableVersion; +import com.facebook.presto.spi.connector.TableFunctionApplicationResult; import com.facebook.presto.spi.constraints.TableConstraint; import com.facebook.presto.spi.function.SqlFunction; import com.facebook.presto.spi.plan.PartitioningHandle; @@ -52,6 +54,7 @@ import com.facebook.presto.spi.statistics.ComputedStatistics; import com.facebook.presto.spi.statistics.TableStatistics; import com.facebook.presto.spi.statistics.TableStatisticsMetadata; +import com.google.common.collect.ImmutableMap; import com.google.common.util.concurrent.ListenableFuture; import io.airlift.slice.Slice; @@ -360,6 +363,11 @@ public interface Metadata */ Map getCatalogNames(Session session); + default Map getCatalogNamesWithConnectorContext(Session session) + { + return ImmutableMap.of(); + } + /** * Get the names that match the specified table prefix (never null). */ @@ -540,4 +548,6 @@ default boolean isPushdownSupportedForFilter(Session session, TableHandle tableH } String normalizeIdentifier(Session session, String catalogName, String identifier); + + Optional> applyTableFunction(Session session, TableFunctionHandle handle); } diff --git a/presto-main-base/src/main/java/com/facebook/presto/metadata/MetadataListing.java b/presto-main-base/src/main/java/com/facebook/presto/metadata/MetadataListing.java index ac0e661bdc04c..c5a2532466a53 100644 --- a/presto-main-base/src/main/java/com/facebook/presto/metadata/MetadataListing.java +++ b/presto-main-base/src/main/java/com/facebook/presto/metadata/MetadataListing.java @@ -15,6 +15,7 @@ import com.facebook.presto.Session; import com.facebook.presto.common.QualifiedObjectName; +import com.facebook.presto.metadata.Catalog.CatalogContext; import com.facebook.presto.spi.ColumnMetadata; import com.facebook.presto.spi.ConnectorId; import com.facebook.presto.spi.SchemaTableName; @@ -54,6 +55,20 @@ public static SortedMap listCatalogs(Session session, Metad return result.build(); } + public static SortedMap listCatalogsWithConnectorContext(Session session, Metadata metadata, AccessControl accessControl) + { + Map catalogNamesWithConnectorContext = metadata.getCatalogNamesWithConnectorContext(session); + Set allowedCatalogs = accessControl.filterCatalogs(session.getIdentity(), session.getAccessControlContext(), catalogNamesWithConnectorContext.keySet()); + + ImmutableSortedMap.Builder result = ImmutableSortedMap.naturalOrder(); + for (Map.Entry entry : catalogNamesWithConnectorContext.entrySet()) { + if (allowedCatalogs.contains(entry.getKey())) { + result.put(entry); + } + } + return result.build(); + } + public static SortedSet listSchemas(Session session, Metadata metadata, AccessControl accessControl, String catalogName) { Set schemaNames = ImmutableSet.copyOf(metadata.listSchemaNames(session, catalogName)); diff --git a/presto-main-base/src/main/java/com/facebook/presto/metadata/MetadataManager.java b/presto-main-base/src/main/java/com/facebook/presto/metadata/MetadataManager.java index a8be73143072a..6e714c5013f3e 100644 --- a/presto-main-base/src/main/java/com/facebook/presto/metadata/MetadataManager.java +++ b/presto-main-base/src/main/java/com/facebook/presto/metadata/MetadataManager.java @@ -26,6 +26,7 @@ import com.facebook.presto.common.predicate.TupleDomain; import com.facebook.presto.common.type.Type; import com.facebook.presto.common.type.TypeSignature; +import com.facebook.presto.metadata.Catalog.CatalogContext; import com.facebook.presto.spi.ColumnHandle; import com.facebook.presto.spi.ColumnMetadata; import com.facebook.presto.spi.ConnectorDeleteTableHandle; @@ -60,6 +61,7 @@ import com.facebook.presto.spi.connector.ConnectorPartitioningMetadata; import com.facebook.presto.spi.connector.ConnectorTableVersion; import com.facebook.presto.spi.connector.ConnectorTransactionHandle; +import com.facebook.presto.spi.connector.TableFunctionApplicationResult; import com.facebook.presto.spi.constraints.TableConstraint; import com.facebook.presto.spi.function.SqlFunction; import com.facebook.presto.spi.plan.PartitioningHandle; @@ -987,6 +989,12 @@ public Map getCatalogNames(Session session) return transactionManager.getCatalogNames(session.getRequiredTransactionId()); } + @Override + public Map getCatalogNamesWithConnectorContext(Session session) + { + return transactionManager.getCatalogNamesWithConnectorContext(session.getRequiredTransactionId()); + } + @Override public List listViews(Session session, QualifiedTablePrefix prefix) { @@ -1532,6 +1540,18 @@ private ColumnMetadata normalizedColumnMetadata(Session session, String catalogN .build(); } + @Override + public Optional> applyTableFunction(Session session, TableFunctionHandle handle) + { + ConnectorId connectorId = handle.getConnectorId(); + ConnectorMetadata metadata = getMetadata(session, connectorId); + + return metadata.applyTableFunction(session.toConnectorSession(connectorId), handle.getFunctionHandle()) + .map(result -> new TableFunctionApplicationResult<>( + new TableHandle(connectorId, result.getTableHandle(), handle.getTransactionHandle(), Optional.empty()), + result.getColumnHandles())); + } + private ViewDefinition deserializeView(String data) { try { diff --git a/presto-main-base/src/main/java/com/facebook/presto/metadata/RemoteNodeStats.java b/presto-main-base/src/main/java/com/facebook/presto/metadata/RemoteNodeStats.java new file mode 100644 index 0000000000000..a9645cd524491 --- /dev/null +++ b/presto-main-base/src/main/java/com/facebook/presto/metadata/RemoteNodeStats.java @@ -0,0 +1,56 @@ +/* + * Licensed 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 com.facebook.presto.metadata; + +import com.facebook.presto.spi.NodeStats; + +import java.util.Optional; + +/** + * Interface for retrieving statistics from remote nodes in a Presto cluster. + *

+ * This interface provides a mechanism to asynchronously fetch and cache node statistics + * from remote Presto worker nodes. Implementations handle the communication protocol + * details (HTTP, Thrift, etc.) and provide a unified way to access node health and + * performance metrics. + *

+ * The interface supports lazy loading and caching of node statistics to minimize + * network overhead while ensuring that cluster management components have access + * to current node state information for scheduling and health monitoring decisions. + */ +public interface RemoteNodeStats +{ + /** + * Returns the cached node statistics if available. + *

+ * This method returns the most recently fetched statistics for the remote node. + * If no statistics have been fetched yet or if the last fetch failed, this + * method returns an empty Optional. + * + * @return an Optional containing the node statistics if available, empty otherwise + */ + Optional getNodeStats(); + + /** + * Triggers an asynchronous refresh of the node statistics. + *

+ * This method initiates a background request to fetch the latest statistics + * from the remote node. The operation is non-blocking and the results will + * be available through subsequent calls to {@link #getNodeStats()}. + *

+ * Implementations should handle network failures gracefully and avoid + * overwhelming the remote node with excessive requests. + */ + void asyncRefresh(); +} diff --git a/presto-main-base/src/main/java/com/facebook/presto/metadata/SignatureBinder.java b/presto-main-base/src/main/java/com/facebook/presto/metadata/SignatureBinder.java index a71844c54b7d4..bd383aa6ef87b 100644 --- a/presto-main-base/src/main/java/com/facebook/presto/metadata/SignatureBinder.java +++ b/presto-main-base/src/main/java/com/facebook/presto/metadata/SignatureBinder.java @@ -13,7 +13,6 @@ */ package com.facebook.presto.metadata; -import com.facebook.presto.UnknownTypeException; import com.facebook.presto.common.type.FunctionType; import com.facebook.presto.common.type.NamedTypeSignature; import com.facebook.presto.common.type.ParameterKind; @@ -25,6 +24,7 @@ import com.facebook.presto.spi.function.LongVariableConstraint; import com.facebook.presto.spi.function.Signature; import com.facebook.presto.spi.function.TypeVariableConstraint; +import com.facebook.presto.spi.type.UnknownTypeException; import com.facebook.presto.sql.analyzer.TypeSignatureProvider; import com.google.common.base.VerifyException; import com.google.common.collect.ImmutableList; diff --git a/presto-main-base/src/main/java/com/facebook/presto/metadata/TableFunctionHandle.java b/presto-main-base/src/main/java/com/facebook/presto/metadata/TableFunctionHandle.java new file mode 100644 index 0000000000000..20e65066d4320 --- /dev/null +++ b/presto-main-base/src/main/java/com/facebook/presto/metadata/TableFunctionHandle.java @@ -0,0 +1,58 @@ +/* + * Licensed 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 com.facebook.presto.metadata; + +import com.facebook.presto.spi.ConnectorId; +import com.facebook.presto.spi.connector.ConnectorTransactionHandle; +import com.facebook.presto.spi.function.table.ConnectorTableFunctionHandle; +import com.fasterxml.jackson.annotation.JsonCreator; +import com.fasterxml.jackson.annotation.JsonProperty; + +import static java.util.Objects.requireNonNull; + +public class TableFunctionHandle +{ + private final ConnectorId connectorId; + private final ConnectorTableFunctionHandle functionHandle; + private final ConnectorTransactionHandle transactionHandle; + + @JsonCreator + public TableFunctionHandle( + @JsonProperty("connectorId") ConnectorId connectorId, + @JsonProperty("functionHandle") ConnectorTableFunctionHandle functionHandle, + @JsonProperty("transactionHandle") ConnectorTransactionHandle transactionHandle) + { + this.connectorId = requireNonNull(connectorId, "connectorId is null"); + this.functionHandle = requireNonNull(functionHandle, "functionHandle is null"); + this.transactionHandle = requireNonNull(transactionHandle, "transactionHandle is null"); + } + + @JsonProperty + public ConnectorId getConnectorId() + { + return connectorId; + } + + @JsonProperty + public ConnectorTableFunctionHandle getFunctionHandle() + { + return functionHandle; + } + + @JsonProperty + public ConnectorTransactionHandle getTransactionHandle() + { + return transactionHandle; + } +} diff --git a/presto-main-base/src/main/java/com/facebook/presto/metadata/TableLayout.java b/presto-main-base/src/main/java/com/facebook/presto/metadata/TableLayout.java index 553f5b5f8da1c..19026f586c035 100644 --- a/presto-main-base/src/main/java/com/facebook/presto/metadata/TableLayout.java +++ b/presto-main-base/src/main/java/com/facebook/presto/metadata/TableLayout.java @@ -78,6 +78,11 @@ public List> getLocalProperties() return layout.getLocalProperties(); } + public Optional getUniqueColumn() + { + return layout.getUniqueColumn(); + } + public ConnectorTableLayoutHandle getLayoutHandle() { return layout.getHandle(); diff --git a/presto-main-base/src/main/java/com/facebook/presto/metadata/ThriftRemoteNodeStats.java b/presto-main-base/src/main/java/com/facebook/presto/metadata/ThriftRemoteNodeStats.java new file mode 100644 index 0000000000000..185df3be1c143 --- /dev/null +++ b/presto-main-base/src/main/java/com/facebook/presto/metadata/ThriftRemoteNodeStats.java @@ -0,0 +1,110 @@ +/* + * Licensed 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 com.facebook.presto.metadata; + +import com.facebook.airlift.log.Logger; +import com.facebook.airlift.units.Duration; +import com.facebook.drift.client.DriftClient; +import com.facebook.presto.server.thrift.ThriftServerInfoClient; +import com.facebook.presto.spi.NodeState; +import com.facebook.presto.spi.NodeStats; +import com.google.common.util.concurrent.FutureCallback; +import com.google.common.util.concurrent.Futures; +import com.google.common.util.concurrent.ListenableFuture; +import com.google.errorprone.annotations.ThreadSafe; +import jakarta.annotation.Nullable; + +import java.net.URI; +import java.util.Optional; +import java.util.concurrent.atomic.AtomicBoolean; +import java.util.concurrent.atomic.AtomicLong; +import java.util.concurrent.atomic.AtomicReference; + +import static com.facebook.airlift.units.Duration.nanosSince; +import static com.google.common.base.Preconditions.checkArgument; +import static com.google.common.util.concurrent.MoreExecutors.directExecutor; +import static java.util.Objects.requireNonNull; +import static java.util.concurrent.TimeUnit.SECONDS; + +@ThreadSafe +public class ThriftRemoteNodeStats + implements RemoteNodeStats +{ + private static final Logger log = Logger.get(ThriftRemoteNodeStats.class); + + private final ThriftServerInfoClient thriftClient; + private final long refreshIntervalMillis; + private final AtomicReference> nodeStats = new AtomicReference<>(Optional.empty()); + private final AtomicBoolean requestInflight = new AtomicBoolean(); + private final AtomicLong lastUpdateNanos = new AtomicLong(); + private final URI stateInfoUri; + private final AtomicLong lastWarningLogged = new AtomicLong(); + + public ThriftRemoteNodeStats(DriftClient thriftClient, URI stateInfoUri, long refreshIntervalMillis) + { + requireNonNull(stateInfoUri, "stateInfoUri is null"); + checkArgument(stateInfoUri.getScheme().equals("thrift"), "unexpected scheme %s", stateInfoUri.getScheme()); + + this.stateInfoUri = stateInfoUri; + this.refreshIntervalMillis = refreshIntervalMillis; + this.thriftClient = requireNonNull(thriftClient, "thriftClient is null").get(Optional.of(stateInfoUri.getAuthority())); + } + + @Override + public Optional getNodeStats() + { + return nodeStats.get(); + } + + @Override + public void asyncRefresh() + { + Duration sinceUpdate = nanosSince(lastUpdateNanos.get()); + if (nanosSince(lastWarningLogged.get()).toMillis() > 1_000 && + sinceUpdate.toMillis() > 10_000 && + requestInflight.get()) { + log.warn("Node state update request to %s has not returned in %s", stateInfoUri, sinceUpdate.toString(SECONDS)); + lastWarningLogged.set(System.nanoTime()); + } + + if (sinceUpdate.toMillis() > refreshIntervalMillis && requestInflight.compareAndSet(false, true)) { + ListenableFuture responseFuture = thriftClient.getServerState(); + + Futures.addCallback(responseFuture, new FutureCallback() + { + @Override + public void onSuccess(@Nullable Integer result) + { + lastUpdateNanos.set(System.nanoTime()); + requestInflight.compareAndSet(true, false); + if (result != null) { + NodeStats nodeStats1 = new NodeStats(NodeState.valueOf(result), null); + nodeStats.set(Optional.of(nodeStats1)); + } + else { + log.warn("Node statistics endpoint %s returned null response, using cached statistics", stateInfoUri); + } + } + + @Override + public void onFailure(Throwable t) + { + log.error("Error fetching node stats from %s: %s", stateInfoUri, t.getMessage()); + lastUpdateNanos.set(System.nanoTime()); + requestInflight.compareAndSet(true, false); + } + }, directExecutor()); + } + } +} diff --git a/presto-main-base/src/main/java/com/facebook/presto/operator/TableWriterOperator.java b/presto-main-base/src/main/java/com/facebook/presto/operator/TableWriterOperator.java index d46627202c6ec..5deda5f3ffb94 100644 --- a/presto-main-base/src/main/java/com/facebook/presto/operator/TableWriterOperator.java +++ b/presto-main-base/src/main/java/com/facebook/presto/operator/TableWriterOperator.java @@ -20,6 +20,7 @@ import com.facebook.drift.annotations.ThriftStruct; import com.facebook.presto.Session; import com.facebook.presto.common.Page; +import com.facebook.presto.common.RuntimeStats; import com.facebook.presto.common.block.Block; import com.facebook.presto.common.block.BlockBuilder; import com.facebook.presto.common.block.RunLengthEncodedBlock; @@ -130,7 +131,7 @@ public Operator createOperator(DriverContext driverContext) boolean statisticsCpuTimerEnabled = !(statisticsAggregationOperator instanceof DevNullOperator) && isStatisticsCpuTimerEnabled(session); return new TableWriterOperator( context, - createPageSink(), + createPageSink(context), columnChannels, notNullChannelColumnNames, statisticsAggregationOperator, @@ -140,19 +141,21 @@ public Operator createOperator(DriverContext driverContext) pageSinkCommitStrategy); } - private ConnectorPageSink createPageSink() + private ConnectorPageSink createPageSink(OperatorContext operatorContext) { PageSinkContext.Builder pageSinkContextBuilder = PageSinkContext.builder() .setCommitRequired(pageSinkCommitStrategy.isCommitRequired()); + RuntimeStats runtimeStats = operatorContext.getRuntimeStats(); + if (target instanceof CreateHandle) { - return pageSinkManager.createPageSink(session, ((CreateHandle) target).getHandle(), pageSinkContextBuilder.build()); + return pageSinkManager.createPageSink(session, ((CreateHandle) target).getHandle(), pageSinkContextBuilder.build(), runtimeStats); } if (target instanceof InsertHandle) { - return pageSinkManager.createPageSink(session, ((InsertHandle) target).getHandle(), pageSinkContextBuilder.build()); + return pageSinkManager.createPageSink(session, ((InsertHandle) target).getHandle(), pageSinkContextBuilder.build(), runtimeStats); } if (target instanceof RefreshMaterializedViewHandle) { - return pageSinkManager.createPageSink(session, ((RefreshMaterializedViewHandle) target).getHandle(), pageSinkContextBuilder.build()); + return pageSinkManager.createPageSink(session, ((RefreshMaterializedViewHandle) target).getHandle(), pageSinkContextBuilder.build(), runtimeStats); } throw new UnsupportedOperationException("Unhandled target type: " + target.getClass().getName()); } diff --git a/presto-main-base/src/main/java/com/facebook/presto/operator/aggregation/AggregationUtils.java b/presto-main-base/src/main/java/com/facebook/presto/operator/aggregation/AggregationUtils.java index 578e782bc8d91..c78186ebb2417 100644 --- a/presto-main-base/src/main/java/com/facebook/presto/operator/aggregation/AggregationUtils.java +++ b/presto-main-base/src/main/java/com/facebook/presto/operator/aggregation/AggregationUtils.java @@ -22,6 +22,7 @@ import com.facebook.presto.operator.aggregation.state.CentralMomentsState; import com.facebook.presto.operator.aggregation.state.CorrelationState; import com.facebook.presto.operator.aggregation.state.CovarianceState; +import com.facebook.presto.operator.aggregation.state.ExtendedRegressionState; import com.facebook.presto.operator.aggregation.state.RegressionState; import com.facebook.presto.operator.aggregation.state.VarianceState; import com.facebook.presto.spi.function.AggregationFunctionImplementation; @@ -145,9 +146,14 @@ public static double getCorrelation(CorrelationState state) public static void updateRegressionState(RegressionState state, double x, double y) { double oldMeanX = state.getMeanX(); - double oldMeanY = state.getMeanY(); updateCovarianceState(state, x, y); state.setM2X(state.getM2X() + (x - oldMeanX) * (x - state.getMeanX())); + } + + public static void updateExtendedRegressionState(ExtendedRegressionState state, double x, double y) + { + double oldMeanY = state.getMeanY(); + updateRegressionState(state, x, y); state.setM2Y(state.getM2Y() + (y - oldMeanY) * (y - state.getMeanY())); } @@ -189,12 +195,12 @@ public static double getRegressionSxy(RegressionState state) return state.getC2(); } - public static double getRegressionSyy(RegressionState state) + public static double getRegressionSyy(ExtendedRegressionState state) { return state.getM2Y(); } - public static double getRegressionR2(RegressionState state) + public static double getRegressionR2(ExtendedRegressionState state) { if (state.getM2X() != 0 && state.getM2Y() == 0) { return 1.0; @@ -311,10 +317,21 @@ public static void mergeRegressionState(RegressionState state, RegressionState o long na = state.getCount(); long nb = otherState.getCount(); state.setM2X(state.getM2X() + otherState.getM2X() + na * nb * Math.pow(state.getMeanX() - otherState.getMeanX(), 2) / (double) (na + nb)); - state.setM2Y(state.getM2Y() + otherState.getM2Y() + na * nb * Math.pow(state.getMeanY() - otherState.getMeanY(), 2) / (double) (na + nb)); updateCovarianceState(state, otherState); } + public static void mergeExtendedRegressionState(ExtendedRegressionState state, ExtendedRegressionState otherState) + { + if (otherState.getCount() == 0) { + return; + } + + long na = state.getCount(); + long nb = otherState.getCount(); + state.setM2Y(state.getM2Y() + otherState.getM2Y() + na * nb * Math.pow(state.getMeanY() - otherState.getMeanY(), 2) / (double) (na + nb)); + mergeRegressionState(state, otherState); + } + public static String generateAggregationName(String baseName, TypeSignature outputType, List inputTypes) { StringBuilder sb = new StringBuilder(); diff --git a/presto-main-base/src/main/java/com/facebook/presto/operator/aggregation/DoubleRegressionAggregation.java b/presto-main-base/src/main/java/com/facebook/presto/operator/aggregation/DoubleRegressionAggregation.java index 24d1c6e61fcf5..db3ad26ec5d6d 100644 --- a/presto-main-base/src/main/java/com/facebook/presto/operator/aggregation/DoubleRegressionAggregation.java +++ b/presto-main-base/src/main/java/com/facebook/presto/operator/aggregation/DoubleRegressionAggregation.java @@ -24,15 +24,8 @@ import com.facebook.presto.spi.function.SqlType; import static com.facebook.presto.common.type.DoubleType.DOUBLE; -import static com.facebook.presto.operator.aggregation.AggregationUtils.getRegressionAvgx; -import static com.facebook.presto.operator.aggregation.AggregationUtils.getRegressionAvgy; -import static com.facebook.presto.operator.aggregation.AggregationUtils.getRegressionCount; import static com.facebook.presto.operator.aggregation.AggregationUtils.getRegressionIntercept; -import static com.facebook.presto.operator.aggregation.AggregationUtils.getRegressionR2; import static com.facebook.presto.operator.aggregation.AggregationUtils.getRegressionSlope; -import static com.facebook.presto.operator.aggregation.AggregationUtils.getRegressionSxx; -import static com.facebook.presto.operator.aggregation.AggregationUtils.getRegressionSxy; -import static com.facebook.presto.operator.aggregation.AggregationUtils.getRegressionSyy; import static com.facebook.presto.operator.aggregation.AggregationUtils.mergeRegressionState; import static com.facebook.presto.operator.aggregation.AggregationUtils.updateRegressionState; @@ -78,100 +71,4 @@ public static void regrIntercept(@AggregationState RegressionState state, BlockB out.appendNull(); } } - - @AggregationFunction("regr_sxy") - @OutputFunction(StandardTypes.DOUBLE) - public static void regrSxy(@AggregationState RegressionState state, BlockBuilder out) - { - double result = getRegressionSxy(state); - double count = getRegressionCount(state); - if (Double.isFinite(result) && Double.isFinite(count) && count > 0) { - DOUBLE.writeDouble(out, result); - } - else { - out.appendNull(); - } - } - - @AggregationFunction("regr_sxx") - @OutputFunction(StandardTypes.DOUBLE) - public static void regrSxx(@AggregationState RegressionState state, BlockBuilder out) - { - double result = getRegressionSxx(state); - double count = getRegressionCount(state); - if (Double.isFinite(result) && Double.isFinite(count) && count > 0) { - DOUBLE.writeDouble(out, result); - } - else { - out.appendNull(); - } - } - - @AggregationFunction("regr_syy") - @OutputFunction(StandardTypes.DOUBLE) - public static void regrSyy(@AggregationState RegressionState state, BlockBuilder out) - { - double result = getRegressionSyy(state); - double count = getRegressionCount(state); - if (Double.isFinite(result) && Double.isFinite(count) && count > 0) { - DOUBLE.writeDouble(out, result); - } - else { - out.appendNull(); - } - } - - @AggregationFunction("regr_r2") - @OutputFunction(StandardTypes.DOUBLE) - public static void regrR2(@AggregationState RegressionState state, BlockBuilder out) - { - double result = getRegressionR2(state); - if (Double.isFinite(result)) { - DOUBLE.writeDouble(out, result); - } - else { - out.appendNull(); - } - } - - @AggregationFunction("regr_count") - @OutputFunction(StandardTypes.DOUBLE) - public static void regrCount(@AggregationState RegressionState state, BlockBuilder out) - { - double result = getRegressionCount(state); - if (Double.isFinite(result) && result > 0) { - DOUBLE.writeDouble(out, result); - } - else { - out.appendNull(); - } - } - - @AggregationFunction("regr_avgy") - @OutputFunction(StandardTypes.DOUBLE) - public static void regrAvgy(@AggregationState RegressionState state, BlockBuilder out) - { - double result = getRegressionAvgy(state); - double count = getRegressionCount(state); - if (Double.isFinite(result) && Double.isFinite(count) && count > 0) { - DOUBLE.writeDouble(out, result); - } - else { - out.appendNull(); - } - } - - @AggregationFunction("regr_avgx") - @OutputFunction(StandardTypes.DOUBLE) - public static void regrAvgx(@AggregationState RegressionState state, BlockBuilder out) - { - double result = getRegressionAvgx(state); - double count = getRegressionCount(state); - if (Double.isFinite(result) && Double.isFinite(count) && count > 0) { - DOUBLE.writeDouble(out, result); - } - else { - out.appendNull(); - } - } } diff --git a/presto-main-base/src/main/java/com/facebook/presto/operator/aggregation/DoubleRegressionExtendedAggregation.java b/presto-main-base/src/main/java/com/facebook/presto/operator/aggregation/DoubleRegressionExtendedAggregation.java new file mode 100644 index 0000000000000..3550cd0936949 --- /dev/null +++ b/presto-main-base/src/main/java/com/facebook/presto/operator/aggregation/DoubleRegressionExtendedAggregation.java @@ -0,0 +1,149 @@ +/* + * Licensed 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 com.facebook.presto.operator.aggregation; + +import com.facebook.presto.common.block.BlockBuilder; +import com.facebook.presto.common.type.StandardTypes; +import com.facebook.presto.operator.aggregation.state.ExtendedRegressionState; +import com.facebook.presto.spi.function.AggregationFunction; +import com.facebook.presto.spi.function.AggregationState; +import com.facebook.presto.spi.function.CombineFunction; +import com.facebook.presto.spi.function.InputFunction; +import com.facebook.presto.spi.function.OutputFunction; +import com.facebook.presto.spi.function.SqlType; + +import static com.facebook.presto.common.type.DoubleType.DOUBLE; +import static com.facebook.presto.operator.aggregation.AggregationUtils.getRegressionAvgx; +import static com.facebook.presto.operator.aggregation.AggregationUtils.getRegressionAvgy; +import static com.facebook.presto.operator.aggregation.AggregationUtils.getRegressionCount; +import static com.facebook.presto.operator.aggregation.AggregationUtils.getRegressionR2; +import static com.facebook.presto.operator.aggregation.AggregationUtils.getRegressionSxx; +import static com.facebook.presto.operator.aggregation.AggregationUtils.getRegressionSxy; +import static com.facebook.presto.operator.aggregation.AggregationUtils.getRegressionSyy; +import static com.facebook.presto.operator.aggregation.AggregationUtils.mergeExtendedRegressionState; +import static com.facebook.presto.operator.aggregation.AggregationUtils.updateExtendedRegressionState; + +@AggregationFunction +public class DoubleRegressionExtendedAggregation +{ + private DoubleRegressionExtendedAggregation() {} + + @InputFunction + public static void input(@AggregationState ExtendedRegressionState state, @SqlType(StandardTypes.DOUBLE) double dependentValue, @SqlType(StandardTypes.DOUBLE) double independentValue) + { + updateExtendedRegressionState(state, independentValue, dependentValue); + } + + @CombineFunction + public static void combine(@AggregationState ExtendedRegressionState state, @AggregationState ExtendedRegressionState otherState) + { + mergeExtendedRegressionState(state, otherState); + } + + @AggregationFunction("regr_sxy") + @OutputFunction(StandardTypes.DOUBLE) + public static void regrSxy(@AggregationState ExtendedRegressionState state, BlockBuilder out) + { + double result = getRegressionSxy(state); + double count = getRegressionCount(state); + if (Double.isFinite(result) && Double.isFinite(count) && count > 0) { + DOUBLE.writeDouble(out, result); + } + else { + out.appendNull(); + } + } + + @AggregationFunction("regr_sxx") + @OutputFunction(StandardTypes.DOUBLE) + public static void regrSxx(@AggregationState ExtendedRegressionState state, BlockBuilder out) + { + double result = getRegressionSxx(state); + double count = getRegressionCount(state); + if (Double.isFinite(result) && Double.isFinite(count) && count > 0) { + DOUBLE.writeDouble(out, result); + } + else { + out.appendNull(); + } + } + + @AggregationFunction("regr_syy") + @OutputFunction(StandardTypes.DOUBLE) + public static void regrSyy(@AggregationState ExtendedRegressionState state, BlockBuilder out) + { + double result = getRegressionSyy(state); + double count = getRegressionCount(state); + if (Double.isFinite(result) && Double.isFinite(count) && count > 0) { + DOUBLE.writeDouble(out, result); + } + else { + out.appendNull(); + } + } + + @AggregationFunction("regr_r2") + @OutputFunction(StandardTypes.DOUBLE) + public static void regrR2(@AggregationState ExtendedRegressionState state, BlockBuilder out) + { + double result = getRegressionR2(state); + if (Double.isFinite(result)) { + DOUBLE.writeDouble(out, result); + } + else { + out.appendNull(); + } + } + + @AggregationFunction("regr_count") + @OutputFunction(StandardTypes.DOUBLE) + public static void regrCount(@AggregationState ExtendedRegressionState state, BlockBuilder out) + { + double result = getRegressionCount(state); + if (Double.isFinite(result) && result > 0) { + DOUBLE.writeDouble(out, result); + } + else { + out.appendNull(); + } + } + + @AggregationFunction("regr_avgy") + @OutputFunction(StandardTypes.DOUBLE) + public static void regrAvgy(@AggregationState ExtendedRegressionState state, BlockBuilder out) + { + double result = getRegressionAvgy(state); + double count = getRegressionCount(state); + if (Double.isFinite(result) && Double.isFinite(count) && count > 0) { + DOUBLE.writeDouble(out, result); + } + else { + out.appendNull(); + } + } + + @AggregationFunction("regr_avgx") + @OutputFunction(StandardTypes.DOUBLE) + public static void regrAvgx(@AggregationState ExtendedRegressionState state, BlockBuilder out) + { + double result = getRegressionAvgx(state); + double count = getRegressionCount(state); + if (Double.isFinite(result) && Double.isFinite(count) && count > 0) { + DOUBLE.writeDouble(out, result); + } + else { + out.appendNull(); + } + } +} diff --git a/presto-main-base/src/main/java/com/facebook/presto/operator/aggregation/RealRegressionAggregation.java b/presto-main-base/src/main/java/com/facebook/presto/operator/aggregation/RealRegressionAggregation.java index 1fe5d006da1a9..a75222bfa93c4 100644 --- a/presto-main-base/src/main/java/com/facebook/presto/operator/aggregation/RealRegressionAggregation.java +++ b/presto-main-base/src/main/java/com/facebook/presto/operator/aggregation/RealRegressionAggregation.java @@ -24,15 +24,8 @@ import com.facebook.presto.spi.function.SqlType; import static com.facebook.presto.common.type.RealType.REAL; -import static com.facebook.presto.operator.aggregation.AggregationUtils.getRegressionAvgx; -import static com.facebook.presto.operator.aggregation.AggregationUtils.getRegressionAvgy; -import static com.facebook.presto.operator.aggregation.AggregationUtils.getRegressionCount; import static com.facebook.presto.operator.aggregation.AggregationUtils.getRegressionIntercept; -import static com.facebook.presto.operator.aggregation.AggregationUtils.getRegressionR2; import static com.facebook.presto.operator.aggregation.AggregationUtils.getRegressionSlope; -import static com.facebook.presto.operator.aggregation.AggregationUtils.getRegressionSxx; -import static com.facebook.presto.operator.aggregation.AggregationUtils.getRegressionSxy; -import static com.facebook.presto.operator.aggregation.AggregationUtils.getRegressionSyy; import static java.lang.Float.floatToRawIntBits; import static java.lang.Float.intBitsToFloat; @@ -78,100 +71,4 @@ public static void regrIntercept(@AggregationState RegressionState state, BlockB out.appendNull(); } } - - @AggregationFunction("regr_sxy") - @OutputFunction(StandardTypes.REAL) - public static void regrSxy(@AggregationState RegressionState state, BlockBuilder out) - { - double result = getRegressionSxy(state); - double count = getRegressionCount(state); - if (Double.isFinite(result) && Double.isFinite(count) && count > 0) { - REAL.writeLong(out, floatToRawIntBits((float) result)); - } - else { - out.appendNull(); - } - } - - @AggregationFunction("regr_sxx") - @OutputFunction(StandardTypes.REAL) - public static void regrSxx(@AggregationState RegressionState state, BlockBuilder out) - { - double result = getRegressionSxx(state); - double count = getRegressionCount(state); - if (Double.isFinite(result) && Double.isFinite(count) && count > 0) { - REAL.writeLong(out, floatToRawIntBits((float) result)); - } - else { - out.appendNull(); - } - } - - @AggregationFunction("regr_syy") - @OutputFunction(StandardTypes.REAL) - public static void regrSyy(@AggregationState RegressionState state, BlockBuilder out) - { - double result = getRegressionSyy(state); - double count = getRegressionCount(state); - if (Double.isFinite(result) && Double.isFinite(count) && count > 0) { - REAL.writeLong(out, floatToRawIntBits((float) result)); - } - else { - out.appendNull(); - } - } - - @AggregationFunction("regr_r2") - @OutputFunction(StandardTypes.REAL) - public static void regrR2(@AggregationState RegressionState state, BlockBuilder out) - { - double result = getRegressionR2(state); - if (Double.isFinite(result)) { - REAL.writeLong(out, floatToRawIntBits((float) result)); - } - else { - out.appendNull(); - } - } - - @AggregationFunction("regr_count") - @OutputFunction(StandardTypes.REAL) - public static void regrCount(@AggregationState RegressionState state, BlockBuilder out) - { - double result = getRegressionCount(state); - if (Double.isFinite(result) && result > 0) { - REAL.writeLong(out, floatToRawIntBits((float) result)); - } - else { - out.appendNull(); - } - } - - @AggregationFunction("regr_avgy") - @OutputFunction(StandardTypes.REAL) - public static void regrAvgy(@AggregationState RegressionState state, BlockBuilder out) - { - double result = getRegressionAvgy(state); - double count = getRegressionCount(state); - if (Double.isFinite(result) && Double.isFinite(count) && count > 0) { - REAL.writeLong(out, floatToRawIntBits((float) result)); - } - else { - out.appendNull(); - } - } - - @AggregationFunction("regr_avgx") - @OutputFunction(StandardTypes.REAL) - public static void regrAvgx(@AggregationState RegressionState state, BlockBuilder out) - { - double result = getRegressionAvgx(state); - double count = getRegressionCount(state); - if (Double.isFinite(result) && Double.isFinite(count) && count > 0) { - REAL.writeLong(out, floatToRawIntBits((float) result)); - } - else { - out.appendNull(); - } - } } diff --git a/presto-main-base/src/main/java/com/facebook/presto/operator/aggregation/RealRegressionExtendedAggregation.java b/presto-main-base/src/main/java/com/facebook/presto/operator/aggregation/RealRegressionExtendedAggregation.java new file mode 100644 index 0000000000000..2d0335ae9aca6 --- /dev/null +++ b/presto-main-base/src/main/java/com/facebook/presto/operator/aggregation/RealRegressionExtendedAggregation.java @@ -0,0 +1,149 @@ +/* + * Licensed 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 com.facebook.presto.operator.aggregation; + +import com.facebook.presto.common.block.BlockBuilder; +import com.facebook.presto.common.type.StandardTypes; +import com.facebook.presto.operator.aggregation.state.ExtendedRegressionState; +import com.facebook.presto.spi.function.AggregationFunction; +import com.facebook.presto.spi.function.AggregationState; +import com.facebook.presto.spi.function.CombineFunction; +import com.facebook.presto.spi.function.InputFunction; +import com.facebook.presto.spi.function.OutputFunction; +import com.facebook.presto.spi.function.SqlType; + +import static com.facebook.presto.common.type.RealType.REAL; +import static com.facebook.presto.operator.aggregation.AggregationUtils.getRegressionAvgx; +import static com.facebook.presto.operator.aggregation.AggregationUtils.getRegressionAvgy; +import static com.facebook.presto.operator.aggregation.AggregationUtils.getRegressionCount; +import static com.facebook.presto.operator.aggregation.AggregationUtils.getRegressionR2; +import static com.facebook.presto.operator.aggregation.AggregationUtils.getRegressionSxx; +import static com.facebook.presto.operator.aggregation.AggregationUtils.getRegressionSxy; +import static com.facebook.presto.operator.aggregation.AggregationUtils.getRegressionSyy; +import static java.lang.Float.floatToRawIntBits; +import static java.lang.Float.intBitsToFloat; + +@AggregationFunction +public class RealRegressionExtendedAggregation +{ + private RealRegressionExtendedAggregation() {} + + @InputFunction + public static void input(@AggregationState ExtendedRegressionState state, @SqlType(StandardTypes.REAL) long dependentValue, @SqlType(StandardTypes.REAL) long independentValue) + { + DoubleRegressionExtendedAggregation.input(state, intBitsToFloat((int) dependentValue), intBitsToFloat((int) independentValue)); + } + + @CombineFunction + public static void combine(@AggregationState ExtendedRegressionState state, @AggregationState ExtendedRegressionState otherState) + { + DoubleRegressionExtendedAggregation.combine(state, otherState); + } + + @AggregationFunction("regr_sxy") + @OutputFunction(StandardTypes.REAL) + public static void regrSxy(@AggregationState ExtendedRegressionState state, BlockBuilder out) + { + double result = getRegressionSxy(state); + double count = getRegressionCount(state); + if (Double.isFinite(result) && Double.isFinite(count) && count > 0) { + REAL.writeLong(out, floatToRawIntBits((float) result)); + } + else { + out.appendNull(); + } + } + + @AggregationFunction("regr_sxx") + @OutputFunction(StandardTypes.REAL) + public static void regrSxx(@AggregationState ExtendedRegressionState state, BlockBuilder out) + { + double result = getRegressionSxx(state); + double count = getRegressionCount(state); + if (Double.isFinite(result) && Double.isFinite(count) && count > 0) { + REAL.writeLong(out, floatToRawIntBits((float) result)); + } + else { + out.appendNull(); + } + } + + @AggregationFunction("regr_syy") + @OutputFunction(StandardTypes.REAL) + public static void regrSyy(@AggregationState ExtendedRegressionState state, BlockBuilder out) + { + double result = getRegressionSyy(state); + double count = getRegressionCount(state); + if (Double.isFinite(result) && Double.isFinite(count) && count > 0) { + REAL.writeLong(out, floatToRawIntBits((float) result)); + } + else { + out.appendNull(); + } + } + + @AggregationFunction("regr_r2") + @OutputFunction(StandardTypes.REAL) + public static void regrR2(@AggregationState ExtendedRegressionState state, BlockBuilder out) + { + double result = getRegressionR2(state); + if (Double.isFinite(result)) { + REAL.writeLong(out, floatToRawIntBits((float) result)); + } + else { + out.appendNull(); + } + } + + @AggregationFunction("regr_count") + @OutputFunction(StandardTypes.REAL) + public static void regrCount(@AggregationState ExtendedRegressionState state, BlockBuilder out) + { + double result = getRegressionCount(state); + if (Double.isFinite(result) && result > 0) { + REAL.writeLong(out, floatToRawIntBits((float) result)); + } + else { + out.appendNull(); + } + } + + @AggregationFunction("regr_avgy") + @OutputFunction(StandardTypes.REAL) + public static void regrAvgy(@AggregationState ExtendedRegressionState state, BlockBuilder out) + { + double result = getRegressionAvgy(state); + double count = getRegressionCount(state); + if (Double.isFinite(result) && Double.isFinite(count) && count > 0) { + REAL.writeLong(out, floatToRawIntBits((float) result)); + } + else { + out.appendNull(); + } + } + + @AggregationFunction("regr_avgx") + @OutputFunction(StandardTypes.REAL) + public static void regrAvgx(@AggregationState ExtendedRegressionState state, BlockBuilder out) + { + double result = getRegressionAvgx(state); + double count = getRegressionCount(state); + if (Double.isFinite(result) && Double.isFinite(count) && count > 0) { + REAL.writeLong(out, floatToRawIntBits((float) result)); + } + else { + out.appendNull(); + } + } +} diff --git a/presto-main-base/src/main/java/com/facebook/presto/operator/aggregation/state/ExtendedRegressionState.java b/presto-main-base/src/main/java/com/facebook/presto/operator/aggregation/state/ExtendedRegressionState.java new file mode 100644 index 0000000000000..64a9883174158 --- /dev/null +++ b/presto-main-base/src/main/java/com/facebook/presto/operator/aggregation/state/ExtendedRegressionState.java @@ -0,0 +1,22 @@ +/* + * Licensed 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 com.facebook.presto.operator.aggregation.state; + +public interface ExtendedRegressionState + extends RegressionState +{ + double getM2Y(); + + void setM2Y(double value); +} diff --git a/presto-main-base/src/main/java/com/facebook/presto/operator/aggregation/state/RegressionState.java b/presto-main-base/src/main/java/com/facebook/presto/operator/aggregation/state/RegressionState.java index 79837f90c0c11..ae3af6f46dc43 100644 --- a/presto-main-base/src/main/java/com/facebook/presto/operator/aggregation/state/RegressionState.java +++ b/presto-main-base/src/main/java/com/facebook/presto/operator/aggregation/state/RegressionState.java @@ -19,8 +19,4 @@ public interface RegressionState double getM2X(); void setM2X(double value); - - double getM2Y(); - - void setM2Y(double value); } diff --git a/presto-main-base/src/main/java/com/facebook/presto/operator/annotations/ImplementationDependency.java b/presto-main-base/src/main/java/com/facebook/presto/operator/annotations/ImplementationDependency.java index 4f30469473589..06e897574a1d4 100644 --- a/presto-main-base/src/main/java/com/facebook/presto/operator/annotations/ImplementationDependency.java +++ b/presto-main-base/src/main/java/com/facebook/presto/operator/annotations/ImplementationDependency.java @@ -20,9 +20,9 @@ import com.facebook.presto.spi.function.Convention; import com.facebook.presto.spi.function.FunctionDependency; import com.facebook.presto.spi.function.InvocationConvention; +import com.facebook.presto.spi.function.LiteralParameter; import com.facebook.presto.spi.function.OperatorDependency; import com.facebook.presto.spi.function.TypeParameter; -import com.facebook.presto.type.LiteralParameter; import java.lang.annotation.Annotation; import java.lang.reflect.AnnotatedElement; diff --git a/presto-main-base/src/main/java/com/facebook/presto/operator/scalar/AbstractArraySortByKeyFunction.java b/presto-main-base/src/main/java/com/facebook/presto/operator/scalar/AbstractArraySortByKeyFunction.java new file mode 100644 index 0000000000000..30313bca7d438 --- /dev/null +++ b/presto-main-base/src/main/java/com/facebook/presto/operator/scalar/AbstractArraySortByKeyFunction.java @@ -0,0 +1,420 @@ +/* + * Licensed 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 com.facebook.presto.operator.scalar; + +import com.facebook.presto.bytecode.BytecodeBlock; +import com.facebook.presto.bytecode.CallSiteBinder; +import com.facebook.presto.bytecode.ClassDefinition; +import com.facebook.presto.bytecode.MethodDefinition; +import com.facebook.presto.bytecode.Parameter; +import com.facebook.presto.bytecode.Scope; +import com.facebook.presto.bytecode.Variable; +import com.facebook.presto.bytecode.control.IfStatement; +import com.facebook.presto.common.NotSupportedException; +import com.facebook.presto.common.QualifiedObjectName; +import com.facebook.presto.common.block.Block; +import com.facebook.presto.common.block.BlockBuilder; +import com.facebook.presto.common.function.SqlFunctionProperties; +import com.facebook.presto.common.type.Type; +import com.facebook.presto.metadata.BoundVariables; +import com.facebook.presto.metadata.FunctionAndTypeManager; +import com.facebook.presto.metadata.SqlScalarFunction; +import com.facebook.presto.spi.PrestoException; +import com.facebook.presto.spi.function.ComplexTypeFunctionDescriptor; +import com.facebook.presto.spi.function.FunctionKind; +import com.facebook.presto.spi.function.LambdaArgumentDescriptor; +import com.facebook.presto.spi.function.LambdaDescriptor; +import com.facebook.presto.spi.function.Signature; +import com.facebook.presto.spi.function.SqlFunctionVisibility; +import com.facebook.presto.sql.gen.lambda.UnaryFunctionInterface; +import com.google.common.collect.ImmutableList; +import com.google.common.collect.ImmutableMap; +import com.google.common.collect.ImmutableSet; +import com.google.common.primitives.Primitives; +import it.unimi.dsi.fastutil.ints.IntComparator; + +import java.lang.invoke.MethodHandle; +import java.lang.invoke.MethodHandles; +import java.util.Optional; + +import static com.facebook.presto.bytecode.Access.FINAL; +import static com.facebook.presto.bytecode.Access.PUBLIC; +import static com.facebook.presto.bytecode.Access.a; +import static com.facebook.presto.bytecode.Parameter.arg; +import static com.facebook.presto.bytecode.ParameterizedType.type; +import static com.facebook.presto.bytecode.expression.BytecodeExpressions.constantNull; +import static com.facebook.presto.bytecode.expression.BytecodeExpressions.equal; +import static com.facebook.presto.common.type.TypeSignature.parseTypeSignature; +import static com.facebook.presto.common.type.UnknownType.UNKNOWN; +import static com.facebook.presto.metadata.BuiltInTypeAndFunctionNamespaceManager.JAVA_BUILTIN_NAMESPACE; +import static com.facebook.presto.operator.scalar.ScalarFunctionImplementationChoice.ArgumentProperty.functionTypeArgumentProperty; +import static com.facebook.presto.operator.scalar.ScalarFunctionImplementationChoice.ArgumentProperty.valueTypeArgumentProperty; +import static com.facebook.presto.operator.scalar.ScalarFunctionImplementationChoice.NullConvention.RETURN_NULL_ON_NULL; +import static com.facebook.presto.spi.StandardErrorCode.INVALID_FUNCTION_ARGUMENT; +import static com.facebook.presto.spi.StandardErrorCode.NOT_SUPPORTED; +import static com.facebook.presto.spi.function.Signature.typeVariable; +import static com.facebook.presto.sql.gen.SqlTypeBytecodeExpression.constantType; +import static com.facebook.presto.util.CompilerUtils.defineClass; +import static com.facebook.presto.util.CompilerUtils.makeClassName; +import static com.facebook.presto.util.Reflection.methodHandle; +import static it.unimi.dsi.fastutil.ints.IntArrays.quickSort; + +public abstract class AbstractArraySortByKeyFunction + extends SqlScalarFunction +{ + private final ComplexTypeFunctionDescriptor descriptor; + + protected AbstractArraySortByKeyFunction(String functionName) + { + super(new Signature( + QualifiedObjectName.valueOf(JAVA_BUILTIN_NAMESPACE, functionName), + FunctionKind.SCALAR, + ImmutableList.of(typeVariable("T"), typeVariable("K")), + ImmutableList.of(), + parseTypeSignature("array(T)"), + ImmutableList.of(parseTypeSignature("array(T)"), parseTypeSignature("function(T,K)")), + false)); + descriptor = new ComplexTypeFunctionDescriptor( + true, + ImmutableList.of(new LambdaDescriptor(1, ImmutableMap.of(0, new LambdaArgumentDescriptor(0, ComplexTypeFunctionDescriptor::prependAllSubscripts)))), + Optional.of(ImmutableSet.of(0)), + Optional.of(ComplexTypeFunctionDescriptor::clearRequiredSubfields), + getSignature()); + } + + @Override + public SqlFunctionVisibility getVisibility() + { + return SqlFunctionVisibility.PUBLIC; + } + + @Override + public boolean isDeterministic() + { + return true; + } + + @Override + public BuiltInScalarFunctionImplementation specialize(BoundVariables boundVariables, int arity, FunctionAndTypeManager functionAndTypeManager) + { + Type elementType = boundVariables.getTypeVariable("T"); + Type keyType = boundVariables.getTypeVariable("K"); + + // Generate the specialized key extractor instance once + KeyExtractor keyExtractor = generateKeyExtractor(elementType, keyType); + + MethodHandle raw = methodHandle( + AbstractArraySortByKeyFunction.class, + "sortByKey", + AbstractArraySortByKeyFunction.class, + Type.class, + Type.class, + KeyExtractor.class, + SqlFunctionProperties.class, + Block.class, + UnaryFunctionInterface.class); + + MethodHandle bound = MethodHandles.insertArguments(raw, 0, this, elementType, keyType, keyExtractor); + + return new BuiltInScalarFunctionImplementation( + false, + ImmutableList.of( + valueTypeArgumentProperty(RETURN_NULL_ON_NULL), // array parameter + functionTypeArgumentProperty(UnaryFunctionInterface.class)), // keyFunction parameter + bound); + } + + @Override + public ComplexTypeFunctionDescriptor getComplexTypeFunctionDescriptor() + { + return descriptor; + } + + public static Block sortByKey( + AbstractArraySortByKeyFunction function, + Type elementType, + Type keyType, + KeyExtractor keyExtractor, + SqlFunctionProperties properties, + Block array, + UnaryFunctionInterface keyFunction) + { + int arrayLength = array.getPositionCount(); + if (arrayLength < 2) { + return array; + } + + // Create array of indices and extracted keys + int[] indices = new int[arrayLength]; + BlockBuilder keyBlockBuilder = keyType.createBlockBuilder(null, arrayLength); + + // Extract keys for all elements + for (int i = 0; i < arrayLength; i++) { + indices[i] = i; + if (array.isNull(i)) { + keyBlockBuilder.appendNull(); + } + else { + try { + // Use the generated KeyExtractor implementation (direct virtual call) + keyExtractor.extract(properties, array, i, keyFunction, keyBlockBuilder); + } + catch (Throwable t) { + throw new PrestoException(INVALID_FUNCTION_ARGUMENT, String.format("Error applying key function to element at position %d", i), t); + } + } + } + + Block keysBlock = keyBlockBuilder.build(); + + // Sort indices based on extracted keys using Type's compareTo + try { + if (array.mayHaveNull() || keysBlock.mayHaveNull()) { + quickSort(indices, new NullableComparator(array, keysBlock, keyType, function)); + } + else { + quickSort(indices, new NonNullableComparator(keysBlock, keyType, function)); + } + } + catch (NotSupportedException | UnsupportedOperationException e) { + throw new PrestoException(INVALID_FUNCTION_ARGUMENT, "Key type does not support comparison", e); + } + catch (PrestoException e) { + if (e.getErrorCode() == NOT_SUPPORTED.toErrorCode()) { + throw new PrestoException(INVALID_FUNCTION_ARGUMENT, "Key type does not support comparison", e); + } + throw e; + } + + // Build result block with sorted elements + BlockBuilder resultBuilder = elementType.createBlockBuilder(null, arrayLength); + for (int i = 0; i < arrayLength; i++) { + elementType.appendTo(array, indices[i], resultBuilder); + } + + return resultBuilder.build(); + } + + /** + * KeyExtractor is a simple interface implemented by generated classes. + * Implementations must write the extracted key into the provided BlockBuilder + * (or appendNull) for the given position. + */ + public interface KeyExtractor + { + void extract(SqlFunctionProperties properties, Block array, int position, UnaryFunctionInterface keyFunction, BlockBuilder keyBlockBuilder) throws Throwable; + } + + // Generate just the key extraction logic + public static KeyExtractor generateKeyExtractor(Type elementType, Type keyType) + { + CallSiteBinder binder = new CallSiteBinder(); + Class elementJavaType = Primitives.wrap(elementType.getJavaType()); + Class keyJavaType = Primitives.wrap(keyType.getJavaType()); + + String className = "ArraySortKeyExtractorImpl_" + elementType.getTypeSignature() + "_" + keyType.getTypeSignature(); + ClassDefinition definition = new ClassDefinition( + a(PUBLIC, FINAL), + makeClassName(className), + type(Object.class), + type(KeyExtractor.class)); + definition.declareDefaultConstructor(a(PUBLIC)); + + Parameter properties = arg("properties", SqlFunctionProperties.class); + Parameter array = arg("array", Block.class); + Parameter position = arg("position", int.class); + Parameter keyFunction = arg("keyFunction", UnaryFunctionInterface.class); + Parameter keyBlockBuilder = arg("keyBlockBuilder", BlockBuilder.class); + + MethodDefinition method = definition.declareMethod( + a(PUBLIC), + "extract", + type(void.class), + ImmutableList.of(properties, array, position, keyFunction, keyBlockBuilder)); + + BytecodeBlock body = method.getBody(); + Scope scope = method.getScope(); + Variable element = scope.declareVariable(elementJavaType, "element"); + Variable key = scope.declareVariable(keyJavaType, "key"); + + // Load element with correct primitive handling + if (!elementType.equals(UNKNOWN)) { + // generates the correct getLong/getDouble/getBoolean/getSlice/getObject call + body.append(element.set(constantType(binder, elementType).getValue(array, position).cast(elementJavaType))); + } + else { + body.append(element.set(constantNull(elementJavaType))); + } + + body.append(key.set(keyFunction.invoke("apply", Object.class, element.cast(Object.class)).cast(keyJavaType))); + + // Write the key to the block builder + if (!keyType.equals(UNKNOWN)) { + body.append(new IfStatement() + .condition(equal(key, constantNull(keyJavaType))) + .ifTrue(keyBlockBuilder.invoke("appendNull", BlockBuilder.class).pop()) + .ifFalse(constantType(binder, keyType).writeValue(keyBlockBuilder, key.cast(keyType.getJavaType())))); + } + else { + body.append(keyBlockBuilder.invoke("appendNull", BlockBuilder.class).pop()); + } + + body.ret(); + + Class generatedClass = defineClass(definition, Object.class, binder.getBindings(), AbstractArraySortByKeyFunction.class.getClassLoader()); + + try { + // instantiate generated class and cast to KeyExtractor for direct virtual call + return (KeyExtractor) generatedClass.getConstructor().newInstance(); + } + catch (ReflectiveOperationException e) { + throw new RuntimeException("Failed to instantiate generated key extractor", e); + } + } + + // Abstract method to be implemented by subclasses to define comparison direction + protected abstract int compareKeys(Type keyType, Block keysBlock, int leftIndex, int rightIndex); + + private static class NullableComparator + implements IntComparator + { + private final Block array; + private final Block keysBlock; + private final Type keyType; + private final AbstractArraySortByKeyFunction function; + + public NullableComparator(Block array, Block keysBlock, Type keyType, AbstractArraySortByKeyFunction function) + { + this.array = array; + this.keysBlock = keysBlock; + this.keyType = keyType; + this.function = function; + } + + @Override + public int compare(int leftIndex, int rightIndex) + { + boolean leftArrayNull = array.isNull(leftIndex); + boolean rightArrayNull = array.isNull(rightIndex); + + if (leftArrayNull && rightArrayNull) { + return 0; + } + if (leftArrayNull) { + return 1; + } + if (rightArrayNull) { + return -1; + } + + boolean leftKeyNull = keysBlock.isNull(leftIndex); + boolean rightKeyNull = keysBlock.isNull(rightIndex); + + if (leftKeyNull && rightKeyNull) { + return 0; + } + if (leftKeyNull) { + return 1; + } + if (rightKeyNull) { + return -1; + } + + int result = function.compareKeys(keyType, keysBlock, leftIndex, rightIndex); + + // If keys are equal, maintain original order + if (result == 0) { + return Integer.compare(leftIndex, rightIndex); + } + + return result; + } + } + + private static class NonNullableComparator + implements IntComparator + { + private final Block keysBlock; + private final Type keyType; + private final AbstractArraySortByKeyFunction function; + + public NonNullableComparator(Block keysBlock, Type keyType, AbstractArraySortByKeyFunction function) + { + this.keysBlock = keysBlock; + this.keyType = keyType; + this.function = function; + } + + @Override + public int compare(int leftIndex, int rightIndex) + { + int result = function.compareKeys(keyType, keysBlock, leftIndex, rightIndex); + + // If keys are equal, maintain original order + if (result == 0) { + return Integer.compare(leftIndex, rightIndex); + } + + return result; + } + } + + public static class ArraySortByKeyFunction + extends AbstractArraySortByKeyFunction + { + public static final ArraySortByKeyFunction ARRAY_SORT_BY_KEY_FUNCTION = new ArraySortByKeyFunction(); + + private ArraySortByKeyFunction() + { + super("array_sort"); + } + + @Override + public String getDescription() + { + return "Sorts the given array using a lambda function to extract sorting keys. " + + "Null array elements and null keys are placed at the end. " + + "Example: array_sort(ARRAY['apple', 'banana', 'cherry'], x -> length(x))"; + } + + @Override + protected int compareKeys(Type keyType, Block keysBlock, int leftIndex, int rightIndex) + { + return keyType.compareTo(keysBlock, leftIndex, keysBlock, rightIndex); + } + } + + public static class ArraySortDescByKeyFunction + extends AbstractArraySortByKeyFunction + { + public static final ArraySortDescByKeyFunction ARRAY_SORT_DESC_BY_KEY_FUNCTION = new ArraySortDescByKeyFunction(); + + private ArraySortDescByKeyFunction() + { + super("array_sort_desc"); + } + + @Override + public String getDescription() + { + return "Sorts the given array in descending order using a lambda function to extract sorting keys"; + } + + @Override + protected int compareKeys(Type keyType, Block keysBlock, int leftIndex, int rightIndex) + { + return keyType.compareTo(keysBlock, rightIndex, keysBlock, leftIndex); + } + } +} diff --git a/presto-main-base/src/main/java/com/facebook/presto/operator/scalar/ArrayIntersectFunction.java b/presto-main-base/src/main/java/com/facebook/presto/operator/scalar/ArrayIntersectFunction.java index d977420d28339..4ee510a86c467 100644 --- a/presto-main-base/src/main/java/com/facebook/presto/operator/scalar/ArrayIntersectFunction.java +++ b/presto-main-base/src/main/java/com/facebook/presto/operator/scalar/ArrayIntersectFunction.java @@ -19,8 +19,6 @@ import com.facebook.presto.spi.function.Description; import com.facebook.presto.spi.function.OperatorDependency; import com.facebook.presto.spi.function.ScalarFunction; -import com.facebook.presto.spi.function.SqlInvokedScalarFunction; -import com.facebook.presto.spi.function.SqlParameter; import com.facebook.presto.spi.function.SqlType; import com.facebook.presto.spi.function.TypeParameter; @@ -60,14 +58,4 @@ public static Block intersect( return typedSet.getBlock(); } - - @SqlInvokedScalarFunction(value = "array_intersect", deterministic = true, calledOnNullInput = false) - @Description("Intersects elements of all arrays in the given array") - @TypeParameter("T") - @SqlParameter(name = "input", type = "array>") - @SqlType("array") - public static String arrayIntersectArray() - { - return "RETURN reduce(input, IF((cardinality(input) = 0), ARRAY[], input[1]), (s, x) -> array_intersect(s, x), (s) -> s)"; - } } diff --git a/presto-main-base/src/main/java/com/facebook/presto/operator/scalar/CharacterStringCasts.java b/presto-main-base/src/main/java/com/facebook/presto/operator/scalar/CharacterStringCasts.java index 0ecf2917fd5cb..b5198a36699ae 100644 --- a/presto-main-base/src/main/java/com/facebook/presto/operator/scalar/CharacterStringCasts.java +++ b/presto-main-base/src/main/java/com/facebook/presto/operator/scalar/CharacterStringCasts.java @@ -14,10 +14,10 @@ package com.facebook.presto.operator.scalar; import com.facebook.presto.common.function.OperatorType; +import com.facebook.presto.spi.function.LiteralParameter; import com.facebook.presto.spi.function.LiteralParameters; import com.facebook.presto.spi.function.ScalarOperator; import com.facebook.presto.spi.function.SqlType; -import com.facebook.presto.type.LiteralParameter; import com.google.common.collect.ImmutableList; import io.airlift.slice.Slice; import io.airlift.slice.SliceUtf8; diff --git a/presto-main-base/src/main/java/com/facebook/presto/operator/scalar/DateTimeFunctions.java b/presto-main-base/src/main/java/com/facebook/presto/operator/scalar/DateTimeFunctions.java index 43f4de46c1e9b..8db043aab9da2 100644 --- a/presto-main-base/src/main/java/com/facebook/presto/operator/scalar/DateTimeFunctions.java +++ b/presto-main-base/src/main/java/com/facebook/presto/operator/scalar/DateTimeFunctions.java @@ -117,13 +117,12 @@ public static long currentTime(SqlFunctionProperties properties) // and we need to have UTC millis for packDateTimeWithZone long millis = UTC_CHRONOLOGY.millisOfDay().get(properties.getSessionStartTime()); - if (!properties.isLegacyTimestamp()) { - // However, those UTC millis are pointing to the correct UTC timestamp - // Our TIME WITH TIME ZONE representation does use UTC 1970-01-01 representation - // So we have to hack here in order to get valid representation - // of TIME WITH TIME ZONE - millis -= valueToSessionTimeZoneOffsetDiff(properties.getSessionStartTime(), getDateTimeZone(properties.getTimeZoneKey())); - } + // However, those UTC millis are pointing to the correct UTC timestamp + // Our TIME WITH TIME ZONE representation does use UTC 1970-01-01 representation + // So we have to hack here in order to get valid representation + // of TIME WITH TIME ZONE + millis -= valueToSessionTimeZoneOffsetDiff(properties.getSessionStartTime(), getDateTimeZone(properties.getTimeZoneKey())); + try { return packDateTimeWithZone(millis, properties.getTimeZoneKey()); } @@ -144,7 +143,8 @@ public static long currentTime(SqlFunctionProperties properties) public static long localTime(SqlFunctionProperties properties) { if (properties.isLegacyTimestamp()) { - return UTC_CHRONOLOGY.millisOfDay().get(properties.getSessionStartTime()); + long millis = UTC_CHRONOLOGY.millisOfDay().get(properties.getSessionStartTime()); + return millis - valueToSessionTimeZoneOffsetDiff(properties.getSessionStartTime(), getDateTimeZone(properties.getTimeZoneKey())); } ISOChronology localChronology = getChronology(properties.getTimeZoneKey()); return localChronology.millisOfDay().get(properties.getSessionStartTime()); diff --git a/presto-main-base/src/main/java/com/facebook/presto/operator/scalar/JoniRegexpCasts.java b/presto-main-base/src/main/java/com/facebook/presto/operator/scalar/JoniRegexpCasts.java index c0f551099672c..eb6264e1a0a1d 100644 --- a/presto-main-base/src/main/java/com/facebook/presto/operator/scalar/JoniRegexpCasts.java +++ b/presto-main-base/src/main/java/com/facebook/presto/operator/scalar/JoniRegexpCasts.java @@ -15,11 +15,11 @@ import com.facebook.presto.common.function.OperatorType; import com.facebook.presto.spi.PrestoException; +import com.facebook.presto.spi.function.LiteralParameter; import com.facebook.presto.spi.function.LiteralParameters; import com.facebook.presto.spi.function.ScalarOperator; import com.facebook.presto.spi.function.SqlType; import com.facebook.presto.type.JoniRegexpType; -import com.facebook.presto.type.LiteralParameter; import io.airlift.jcodings.specific.NonStrictUTF8Encoding; import io.airlift.joni.Option; import io.airlift.joni.Regex; diff --git a/presto-main-base/src/main/java/com/facebook/presto/operator/scalar/JsonFunctions.java b/presto-main-base/src/main/java/com/facebook/presto/operator/scalar/JsonFunctions.java index 339ff5c10dba9..41242a9fcbf50 100644 --- a/presto-main-base/src/main/java/com/facebook/presto/operator/scalar/JsonFunctions.java +++ b/presto-main-base/src/main/java/com/facebook/presto/operator/scalar/JsonFunctions.java @@ -21,13 +21,13 @@ import com.facebook.presto.common.type.StandardTypes; import com.facebook.presto.common.type.Type; import com.facebook.presto.spi.PrestoException; +import com.facebook.presto.spi.function.LiteralParameter; import com.facebook.presto.spi.function.LiteralParameters; import com.facebook.presto.spi.function.ScalarFunction; import com.facebook.presto.spi.function.ScalarOperator; import com.facebook.presto.spi.function.SqlNullable; import com.facebook.presto.spi.function.SqlType; import com.facebook.presto.type.JsonPathType; -import com.facebook.presto.type.LiteralParameter; import com.fasterxml.jackson.core.JsonFactory; import com.fasterxml.jackson.core.JsonParser; import com.fasterxml.jackson.core.JsonToken; diff --git a/presto-main-base/src/main/java/com/facebook/presto/operator/scalar/MathFunctions.java b/presto-main-base/src/main/java/com/facebook/presto/operator/scalar/MathFunctions.java index e778108d74303..1bdd52292b2f7 100644 --- a/presto-main-base/src/main/java/com/facebook/presto/operator/scalar/MathFunctions.java +++ b/presto-main-base/src/main/java/com/facebook/presto/operator/scalar/MathFunctions.java @@ -21,13 +21,13 @@ import com.facebook.presto.operator.aggregation.TypedSet; import com.facebook.presto.spi.PrestoException; import com.facebook.presto.spi.function.Description; +import com.facebook.presto.spi.function.LiteralParameter; import com.facebook.presto.spi.function.LiteralParameters; import com.facebook.presto.spi.function.ScalarFunction; import com.facebook.presto.spi.function.Signature; import com.facebook.presto.spi.function.SqlNullable; import com.facebook.presto.spi.function.SqlType; import com.facebook.presto.type.Constraint; -import com.facebook.presto.type.LiteralParameter; import com.facebook.presto.util.SecureRandomGeneration; import com.google.common.primitives.Doubles; import io.airlift.slice.Slice; diff --git a/presto-main-base/src/main/java/com/facebook/presto/operator/scalar/StringFunctions.java b/presto-main-base/src/main/java/com/facebook/presto/operator/scalar/StringFunctions.java index 2a55fe9485f29..ba7722722f609 100644 --- a/presto-main-base/src/main/java/com/facebook/presto/operator/scalar/StringFunctions.java +++ b/presto-main-base/src/main/java/com/facebook/presto/operator/scalar/StringFunctions.java @@ -19,6 +19,7 @@ import com.facebook.presto.common.type.StandardTypes; import com.facebook.presto.spi.PrestoException; import com.facebook.presto.spi.function.Description; +import com.facebook.presto.spi.function.LiteralParameter; import com.facebook.presto.spi.function.LiteralParameters; import com.facebook.presto.spi.function.ScalarFunction; import com.facebook.presto.spi.function.ScalarOperator; @@ -26,7 +27,6 @@ import com.facebook.presto.spi.function.SqlType; import com.facebook.presto.type.CodePointsType; import com.facebook.presto.type.Constraint; -import com.facebook.presto.type.LiteralParameter; import com.google.common.primitives.Ints; import io.airlift.slice.InvalidCodePointException; import io.airlift.slice.InvalidUtf8Exception; diff --git a/presto-main-base/src/main/java/com/facebook/presto/server/InternalCommunicationConfig.java b/presto-main-base/src/main/java/com/facebook/presto/server/InternalCommunicationConfig.java index 06990e3057559..47812dc99f384 100644 --- a/presto-main-base/src/main/java/com/facebook/presto/server/InternalCommunicationConfig.java +++ b/presto-main-base/src/main/java/com/facebook/presto/server/InternalCommunicationConfig.java @@ -49,6 +49,8 @@ public class InternalCommunicationConfig private CommunicationProtocol serverInfoCommunicationProtocol = CommunicationProtocol.HTTP; private boolean memoizeDeadNodesEnabled; private String sharedSecret; + private long nodeStatsRefreshIntervalMillis = 1_000; + private long nodeDiscoveryPollingIntervalMillis = 5_000; private boolean internalJwtEnabled; @@ -312,6 +314,32 @@ public InternalCommunicationConfig setSharedSecret(String sharedSecret) return this; } + public long getNodeStatsRefreshIntervalMillis() + { + return nodeStatsRefreshIntervalMillis; + } + + @Config("internal-communication.node-stats-refresh-interval-millis") + @ConfigDescription("Interval in milliseconds for refreshing node statistics") + public InternalCommunicationConfig setNodeStatsRefreshIntervalMillis(long nodeStatsRefreshIntervalMillis) + { + this.nodeStatsRefreshIntervalMillis = nodeStatsRefreshIntervalMillis; + return this; + } + + public long getNodeDiscoveryPollingIntervalMillis() + { + return nodeDiscoveryPollingIntervalMillis; + } + + @Config("internal-communication.node-discovery-polling-interval-millis") + @ConfigDescription("Interval in milliseconds for polling node discovery and refreshing node states") + public InternalCommunicationConfig setNodeDiscoveryPollingIntervalMillis(long nodeDiscoveryPollingIntervalMillis) + { + this.nodeDiscoveryPollingIntervalMillis = nodeDiscoveryPollingIntervalMillis; + return this; + } + public boolean isInternalJwtEnabled() { return internalJwtEnabled; diff --git a/presto-main-base/src/main/java/com/facebook/presto/server/security/SecurityConfig.java b/presto-main-base/src/main/java/com/facebook/presto/server/security/SecurityConfig.java index 9ac1cb678ba29..52d70ce60ac49 100644 --- a/presto-main-base/src/main/java/com/facebook/presto/server/security/SecurityConfig.java +++ b/presto-main-base/src/main/java/com/facebook/presto/server/security/SecurityConfig.java @@ -40,7 +40,9 @@ public enum AuthenticationType KERBEROS, PASSWORD, JWT, - CUSTOM + CUSTOM, + TEST_EXTERNAL, + OAUTH2 } @NotNull @@ -56,7 +58,7 @@ public SecurityConfig setAuthenticationTypes(List authentica } @Config("http-server.authentication.type") - @ConfigDescription("Authentication types (supported types: CERTIFICATE, KERBEROS, PASSWORD, JWT, CUSTOM)") + @ConfigDescription("Authentication types (supported types: CERTIFICATE, KERBEROS, PASSWORD, JWT, CUSTOM, OAUTH2, TEST_EXTERNAL)") public SecurityConfig setAuthenticationTypes(String types) { if (types == null) { diff --git a/presto-main-base/src/main/java/com/facebook/presto/sessionpropertyproviders/NativeWorkerSessionPropertyProvider.java b/presto-main-base/src/main/java/com/facebook/presto/sessionpropertyproviders/NativeWorkerSessionPropertyProvider.java index a4d6c81eb62ba..300f2485b5e97 100644 --- a/presto-main-base/src/main/java/com/facebook/presto/sessionpropertyproviders/NativeWorkerSessionPropertyProvider.java +++ b/presto-main-base/src/main/java/com/facebook/presto/sessionpropertyproviders/NativeWorkerSessionPropertyProvider.java @@ -28,7 +28,6 @@ import static com.facebook.presto.spi.session.PropertyMetadata.stringProperty; import static java.util.Objects.requireNonNull; -@Deprecated public class NativeWorkerSessionPropertyProvider implements WorkerSessionPropertyProvider { @@ -53,6 +52,7 @@ public class NativeWorkerSessionPropertyProvider public static final String NATIVE_DEBUG_DISABLE_EXPRESSION_WITH_MEMOIZATION = "native_debug_disable_expression_with_memoization"; public static final String NATIVE_DEBUG_DISABLE_EXPRESSION_WITH_LAZY_INPUTS = "native_debug_disable_expression_with_lazy_inputs"; public static final String NATIVE_DEBUG_MEMORY_POOL_NAME_REGEX = "native_debug_memory_pool_name_regex"; + public static final String NATIVE_DEBUG_MEMORY_POOL_WARN_THRESHOLD_BYTES = "native_debug_memory_pool_warn_threshold_bytes"; public static final String NATIVE_SELECTIVE_NIMBLE_READER_ENABLED = "native_selective_nimble_reader_enabled"; public static final String NATIVE_MAX_PARTIAL_AGGREGATION_MEMORY = "native_max_partial_aggregation_memory"; public static final String NATIVE_MAX_EXTENDED_PARTIAL_AGGREGATION_MEMORY = "native_max_extended_partial_aggregation_memory"; @@ -80,6 +80,10 @@ public class NativeWorkerSessionPropertyProvider public static final String NATIVE_REQUEST_DATA_SIZES_MAX_WAIT_SEC = "native_request_data_sizes_max_wait_sec"; public static final String NATIVE_QUERY_MEMORY_RECLAIMER_PRIORITY = "native_query_memory_reclaimer_priority"; public static final String NATIVE_MAX_NUM_SPLITS_LISTENED_TO = "native_max_num_splits_listened_to"; + public static final String NATIVE_INDEX_LOOKUP_JOIN_MAX_PREFETCH_BATCHES = "native_index_lookup_join_max_prefetch_batches"; + public static final String NATIVE_INDEX_LOOKUP_JOIN_SPLIT_OUTPUT = "native_index_lookup_join_split_output"; + public static final String NATIVE_UNNEST_SPLIT_OUTPUT = "native_unnest_split_output"; + private final List> sessionProperties; @Inject @@ -151,7 +155,7 @@ public NativeWorkerSessionPropertyProvider(FeaturesConfig featuresConfig) longProperty( NATIVE_WRITER_FLUSH_THRESHOLD_BYTES, "Native Execution only. Minimum memory footprint size required to reclaim memory from a file " + - "writer by flushing its buffered data to disk.", + "writer by flushing its buffered data to disk.", 96L << 20, false), booleanProperty( @@ -213,6 +217,15 @@ public NativeWorkerSessionPropertyProvider(FeaturesConfig featuresConfig) " string means no match for all.", "", true), + stringProperty( + NATIVE_DEBUG_MEMORY_POOL_WARN_THRESHOLD_BYTES, + "Warning threshold in bytes for debug memory pools. When set to a " + + "non-zero value, a warning will be logged once per memory pool when " + + "allocations cause the pool to exceed this threshold. This is useful for " + + "identifying memory usage patterns during debugging. A value of " + + "0 means no warning threshold is enforced.", + "0B", + true), booleanProperty( NATIVE_SELECTIVE_NIMBLE_READER_ENABLED, "Temporary flag to control whether selective Nimble reader should be " + @@ -256,7 +269,7 @@ public NativeWorkerSessionPropertyProvider(FeaturesConfig featuresConfig) "", !nativeExecution), stringProperty(NATIVE_QUERY_TRACE_FRAGMENT_ID, - "The fragment id of the traced task.", + "The fragment id of the traced task.", "", !nativeExecution), stringProperty(NATIVE_QUERY_TRACE_SHARD_ID, @@ -356,13 +369,34 @@ public NativeWorkerSessionPropertyProvider(FeaturesConfig featuresConfig) integerProperty( NATIVE_QUERY_MEMORY_RECLAIMER_PRIORITY, "Native Execution only. Priority of memory recliamer when deciding on memory pool to abort." + - "Lower value has higher priority and less likely to be choosen for memory pool abort", + "Lower value has higher priority and less likely to be choosen for memory pool abort", 2147483647, !nativeExecution), integerProperty( NATIVE_MAX_NUM_SPLITS_LISTENED_TO, "Maximum number of splits to listen to per table scan node per worker.", 0, + !nativeExecution), + + integerProperty( + NATIVE_INDEX_LOOKUP_JOIN_MAX_PREFETCH_BATCHES, + "Specifies the max number of input batches to prefetch to do index lookup ahead. " + + "If it is zero, then process one input batch at a time.", + 0, + !nativeExecution), + booleanProperty( + NATIVE_INDEX_LOOKUP_JOIN_SPLIT_OUTPUT, + "If this is true, then the index join operator might split output for each input " + + "batch based on the output batch size control. Otherwise, it tries to produce a " + + "single output for each input batch.", + true, + !nativeExecution), + booleanProperty( + NATIVE_UNNEST_SPLIT_OUTPUT, + "If this is true, then the unnest operator might split output for each input " + + "batch based on the output batch size control. Otherwise, it produces a single " + + "output for each input batch.", + true, !nativeExecution)); } diff --git a/presto-main-base/src/main/java/com/facebook/presto/split/PageSinkManager.java b/presto-main-base/src/main/java/com/facebook/presto/split/PageSinkManager.java index aee08454aefc9..39f6a860ac1d5 100644 --- a/presto-main-base/src/main/java/com/facebook/presto/split/PageSinkManager.java +++ b/presto-main-base/src/main/java/com/facebook/presto/split/PageSinkManager.java @@ -14,6 +14,7 @@ package com.facebook.presto.split; import com.facebook.presto.Session; +import com.facebook.presto.common.RuntimeStats; import com.facebook.presto.metadata.InsertTableHandle; import com.facebook.presto.metadata.OutputTableHandle; import com.facebook.presto.spi.ConnectorId; @@ -46,22 +47,32 @@ public void removeConnectorPageSinkProvider(ConnectorId connectorId) pageSinkProviders.remove(connectorId); } - @Override - public ConnectorPageSink createPageSink(Session session, OutputTableHandle tableHandle, PageSinkContext pageSinkContext) + public ConnectorPageSink createPageSink(Session session, OutputTableHandle tableHandle, PageSinkContext pageSinkContext, RuntimeStats runtimeStats) { // assumes connectorId and catalog are the same - ConnectorSession connectorSession = session.toConnectorSession(tableHandle.getConnectorId()); + ConnectorSession connectorSession = session.toConnectorSession(tableHandle.getConnectorId(), runtimeStats); return providerFor(tableHandle.getConnectorId()).createPageSink(tableHandle.getTransactionHandle(), connectorSession, tableHandle.getConnectorHandle(), pageSinkContext); } - @Override - public ConnectorPageSink createPageSink(Session session, InsertTableHandle tableHandle, PageSinkContext pageSinkContext) + public ConnectorPageSink createPageSink(Session session, InsertTableHandle tableHandle, PageSinkContext pageSinkContext, RuntimeStats runtimeStats) { // assumes connectorId and catalog are the same - ConnectorSession connectorSession = session.toConnectorSession(tableHandle.getConnectorId()); + ConnectorSession connectorSession = session.toConnectorSession(tableHandle.getConnectorId(), runtimeStats); return providerFor(tableHandle.getConnectorId()).createPageSink(tableHandle.getTransactionHandle(), connectorSession, tableHandle.getConnectorHandle(), pageSinkContext); } + @Override + public ConnectorPageSink createPageSink(Session session, OutputTableHandle tableHandle, PageSinkContext pageSinkContext) + { + return createPageSink(session, tableHandle, pageSinkContext, null); + } + + @Override + public ConnectorPageSink createPageSink(Session session, InsertTableHandle tableHandle, PageSinkContext pageSinkContext) + { + return createPageSink(session, tableHandle, pageSinkContext, null); + } + private ConnectorPageSinkProvider providerFor(ConnectorId connectorId) { ConnectorPageSinkProvider provider = pageSinkProviders.get(connectorId); diff --git a/presto-main-base/src/main/java/com/facebook/presto/sql/TemporaryTableUtil.java b/presto-main-base/src/main/java/com/facebook/presto/sql/TemporaryTableUtil.java index bc2b3d97805e8..dd46653509bfa 100644 --- a/presto-main-base/src/main/java/com/facebook/presto/sql/TemporaryTableUtil.java +++ b/presto-main-base/src/main/java/com/facebook/presto/sql/TemporaryTableUtil.java @@ -202,7 +202,7 @@ public static TableFinishNode createTemporaryTableWriteWithoutExchanges( Optional cteId) { SchemaTableName schemaTableName = metadata.getTableMetadata(session, tableHandle).getTable(); - TableWriterNode.InsertReference insertReference = new TableWriterNode.InsertReference(tableHandle, schemaTableName); + TableWriterNode.InsertReference insertReference = new TableWriterNode.InsertReference(tableHandle, schemaTableName, Optional.empty()); List outputColumnNames = outputs.stream() .map(variableToColumnMap::get) .map(ColumnMetadata::getName) @@ -300,7 +300,7 @@ public static TableFinishNode createTemporaryTableWriteWithExchanges( .collect(Collectors.toSet()); SchemaTableName schemaTableName = metadata.getTableMetadata(session, tableHandle).getTable(); - TableWriterNode.InsertReference insertReference = new TableWriterNode.InsertReference(tableHandle, schemaTableName); + TableWriterNode.InsertReference insertReference = new TableWriterNode.InsertReference(tableHandle, schemaTableName, Optional.empty()); PartitioningScheme partitioningScheme = new PartitioningScheme( Partitioning.create(partitioningHandle, partitioningVariables), diff --git a/presto-main-base/src/main/java/com/facebook/presto/sql/analyzer/ExpressionAnalyzer.java b/presto-main-base/src/main/java/com/facebook/presto/sql/analyzer/ExpressionAnalyzer.java index b7446727feb8b..e2739aca9a69b 100644 --- a/presto-main-base/src/main/java/com/facebook/presto/sql/analyzer/ExpressionAnalyzer.java +++ b/presto-main-base/src/main/java/com/facebook/presto/sql/analyzer/ExpressionAnalyzer.java @@ -14,7 +14,6 @@ package com.facebook.presto.sql.analyzer; import com.facebook.presto.Session; -import com.facebook.presto.UnknownTypeException; import com.facebook.presto.common.ErrorCode; import com.facebook.presto.common.QualifiedObjectName; import com.facebook.presto.common.Subfield; @@ -47,6 +46,7 @@ import com.facebook.presto.spi.function.SqlInvokedFunction; import com.facebook.presto.spi.security.AccessControl; import com.facebook.presto.spi.security.DenyAllAccessControl; +import com.facebook.presto.spi.type.UnknownTypeException; import com.facebook.presto.sql.parser.SqlParser; import com.facebook.presto.sql.planner.TypeProvider; import com.facebook.presto.sql.relational.FunctionResolution; @@ -251,6 +251,8 @@ public class ExpressionAnalyzer // This contains types of variables referenced from outer scopes. private final Map, Type> outerScopeSymbolTypes; + private final List sourceFields = new ArrayList<>(); + private ExpressionAnalyzer( FunctionAndTypeResolver functionAndTypeResolver, Function statementAnalyzerFactory, @@ -382,6 +384,11 @@ public Multimap getTableColumnAndSubfieldReferenc return tableColumnAndSubfieldReferences; } + public List getSourceFields() + { + return sourceFields; + } + public Multimap getTableColumnAndSubfieldReferencesForAccessControl() { return tableColumnAndSubfieldReferencesForAccessControl; @@ -497,6 +504,8 @@ private Type handleResolvedField(Expression node, FieldId fieldId, Field field, } } + sourceFields.add(field); + // If we found a direct column reference, and we will put it in tableColumnReferencesWithSubFields if (isTopMostReference(node, context)) { Optional tableName = field.getOriginTable(); @@ -1452,7 +1461,7 @@ else if (previousNode instanceof QuantifiedComparisonExpression) { else { scalarSubqueries.add(NodeRef.of(node)); } - + sourceFields.add(queryScope.getRelationType().getFieldByIndex(0)); Type type = getOnlyElement(queryScope.getRelationType().getVisibleFields()).getType(); return setExpressionType(node, type); } @@ -1992,6 +2001,8 @@ public static ExpressionAnalysis analyzeExpression( analyzer.getTableColumnAndSubfieldReferences(), analyzer.getTableColumnAndSubfieldReferencesForAccessControl()); + analysis.addExpressionFields(expression, analyzer.getSourceFields()); + return new ExpressionAnalysis( expressionTypes, expressionCoercions, diff --git a/presto-main-base/src/main/java/com/facebook/presto/sql/analyzer/ExpressionTreeUtils.java b/presto-main-base/src/main/java/com/facebook/presto/sql/analyzer/ExpressionTreeUtils.java index 7ad6060e5a90d..048e57ee7138d 100644 --- a/presto-main-base/src/main/java/com/facebook/presto/sql/analyzer/ExpressionTreeUtils.java +++ b/presto-main-base/src/main/java/com/facebook/presto/sql/analyzer/ExpressionTreeUtils.java @@ -13,13 +13,13 @@ */ package com.facebook.presto.sql.analyzer; -import com.facebook.presto.UnknownTypeException; import com.facebook.presto.common.type.EnumType; import com.facebook.presto.common.type.Type; import com.facebook.presto.common.type.TypeWithName; import com.facebook.presto.spi.SourceLocation; import com.facebook.presto.spi.function.FunctionHandle; import com.facebook.presto.spi.relation.VariableReferenceExpression; +import com.facebook.presto.spi.type.UnknownTypeException; import com.facebook.presto.sql.tree.ArrayConstructor; import com.facebook.presto.sql.tree.Cast; import com.facebook.presto.sql.tree.ComparisonExpression; diff --git a/presto-main-base/src/main/java/com/facebook/presto/sql/analyzer/FeaturesConfig.java b/presto-main-base/src/main/java/com/facebook/presto/sql/analyzer/FeaturesConfig.java index 2b0b89d99b472..cdecf697bda5e 100644 --- a/presto-main-base/src/main/java/com/facebook/presto/sql/analyzer/FeaturesConfig.java +++ b/presto-main-base/src/main/java/com/facebook/presto/sql/analyzer/FeaturesConfig.java @@ -309,6 +309,9 @@ public class FeaturesConfig private boolean addDistinctBelowSemiJoinBuild; private boolean pushdownSubfieldForMapFunctions = true; private long maxSerializableObjectSize = 1000; + private boolean utilizeUniquePropertyInQueryPlanning = true; + + private boolean builtInSidecarFunctionsEnabled; public enum PartitioningPrecisionStrategy { @@ -3099,6 +3102,19 @@ public boolean isPushdownSubfieldForMapFunctions() return pushdownSubfieldForMapFunctions; } + @Config("optimizer.utilize-unique-property-in-query-planning") + @ConfigDescription("Utilize the unique property of input columns in query planning") + public FeaturesConfig setUtilizeUniquePropertyInQueryPlanning(boolean utilizeUniquePropertyInQueryPlanning) + { + this.utilizeUniquePropertyInQueryPlanning = utilizeUniquePropertyInQueryPlanning; + return this; + } + + public boolean isUtilizeUniquePropertyInQueryPlanning() + { + return utilizeUniquePropertyInQueryPlanning; + } + @Config("max_serializable_object_size") @ConfigDescription("Configure the maximum byte size of a serializable object in expression interpreters") public FeaturesConfig setMaxSerializableObjectSize(long maxSerializableObjectSize) @@ -3111,4 +3127,17 @@ public long getMaxSerializableObjectSize() { return maxSerializableObjectSize; } + + @Config("built-in-sidecar-functions-enabled") + @ConfigDescription("Enable using CPP functions from sidecar over coordinator SQL implementations.") + public FeaturesConfig setBuiltInSidecarFunctionsEnabled(boolean builtInSidecarFunctionsEnabled) + { + this.builtInSidecarFunctionsEnabled = builtInSidecarFunctionsEnabled; + return this; + } + + public boolean isBuiltInSidecarFunctionsEnabled() + { + return this.builtInSidecarFunctionsEnabled; + } } diff --git a/presto-main-base/src/main/java/com/facebook/presto/sql/analyzer/StatementAnalyzer.java b/presto-main-base/src/main/java/com/facebook/presto/sql/analyzer/StatementAnalyzer.java index b16042d87537b..925d2704fcf32 100644 --- a/presto-main-base/src/main/java/com/facebook/presto/sql/analyzer/StatementAnalyzer.java +++ b/presto-main-base/src/main/java/com/facebook/presto/sql/analyzer/StatementAnalyzer.java @@ -17,6 +17,7 @@ import com.facebook.presto.Session; import com.facebook.presto.SystemSessionProperties; import com.facebook.presto.common.QualifiedObjectName; +import com.facebook.presto.common.SourceColumn; import com.facebook.presto.common.Subfield; import com.facebook.presto.common.function.OperatorType; import com.facebook.presto.common.predicate.Domain; @@ -31,15 +32,19 @@ import com.facebook.presto.common.type.TimestampWithTimeZoneType; import com.facebook.presto.common.type.Type; import com.facebook.presto.common.type.VarcharType; +import com.facebook.presto.metadata.CatalogMetadata; import com.facebook.presto.metadata.Metadata; import com.facebook.presto.metadata.OperatorNotFoundException; +import com.facebook.presto.metadata.TableFunctionMetadata; import com.facebook.presto.spi.ColumnHandle; import com.facebook.presto.spi.ColumnMetadata; +import com.facebook.presto.spi.ConnectorId; import com.facebook.presto.spi.MaterializedViewDefinition; import com.facebook.presto.spi.MaterializedViewStatus; import com.facebook.presto.spi.PrestoException; import com.facebook.presto.spi.PrestoWarning; import com.facebook.presto.spi.SchemaTableName; +import com.facebook.presto.spi.StandardErrorCode; import com.facebook.presto.spi.TableHandle; import com.facebook.presto.spi.WarningCollector; import com.facebook.presto.spi.analyzer.AccessControlInfo; @@ -47,9 +52,22 @@ import com.facebook.presto.spi.analyzer.MetadataResolver; import com.facebook.presto.spi.analyzer.ViewDefinition; import com.facebook.presto.spi.connector.ConnectorTableVersion; +import com.facebook.presto.spi.connector.ConnectorTransactionHandle; +import com.facebook.presto.spi.eventlistener.Column; +import com.facebook.presto.spi.eventlistener.OutputColumnMetadata; import com.facebook.presto.spi.function.FunctionKind; import com.facebook.presto.spi.function.Signature; import com.facebook.presto.spi.function.SqlFunction; +import com.facebook.presto.spi.function.table.Argument; +import com.facebook.presto.spi.function.table.ArgumentSpecification; +import com.facebook.presto.spi.function.table.ConnectorTableFunction; +import com.facebook.presto.spi.function.table.Descriptor; +import com.facebook.presto.spi.function.table.DescriptorArgumentSpecification; +import com.facebook.presto.spi.function.table.ReturnTypeSpecification; +import com.facebook.presto.spi.function.table.ScalarArgument; +import com.facebook.presto.spi.function.table.ScalarArgumentSpecification; +import com.facebook.presto.spi.function.table.TableArgumentSpecification; +import com.facebook.presto.spi.function.table.TableFunctionAnalysis; import com.facebook.presto.spi.relation.DomainTranslator; import com.facebook.presto.spi.relation.RowExpression; import com.facebook.presto.spi.security.AccessControl; @@ -127,6 +145,7 @@ import com.facebook.presto.sql.tree.NodeRef; import com.facebook.presto.sql.tree.Offset; import com.facebook.presto.sql.tree.OrderBy; +import com.facebook.presto.sql.tree.Parameter; import com.facebook.presto.sql.tree.Prepare; import com.facebook.presto.sql.tree.Property; import com.facebook.presto.sql.tree.QualifiedName; @@ -158,6 +177,10 @@ import com.facebook.presto.sql.tree.StartTransaction; import com.facebook.presto.sql.tree.Statement; import com.facebook.presto.sql.tree.Table; +import com.facebook.presto.sql.tree.TableFunctionArgument; +import com.facebook.presto.sql.tree.TableFunctionDescriptorArgument; +import com.facebook.presto.sql.tree.TableFunctionInvocation; +import com.facebook.presto.sql.tree.TableFunctionTableArgument; import com.facebook.presto.sql.tree.TableSubquery; import com.facebook.presto.sql.tree.TruncateTable; import com.facebook.presto.sql.tree.Union; @@ -171,6 +194,7 @@ import com.facebook.presto.sql.tree.With; import com.facebook.presto.sql.tree.WithQuery; import com.facebook.presto.sql.util.AstUtils; +import com.facebook.presto.transaction.TransactionManager; import com.google.common.base.Joiner; import com.google.common.collect.ImmutableList; import com.google.common.collect.ImmutableMap; @@ -178,6 +202,7 @@ import com.google.common.collect.ImmutableSet; import com.google.common.collect.Iterables; import com.google.common.collect.Multimap; +import com.google.common.collect.Streams; import java.util.ArrayList; import java.util.Arrays; @@ -191,6 +216,7 @@ import java.util.Set; import java.util.function.Function; import java.util.stream.Collectors; +import java.util.stream.Stream; import static com.facebook.presto.SystemSessionProperties.getMaxGroupingSets; import static com.facebook.presto.SystemSessionProperties.isAllowWindowOrderByLiterals; @@ -208,9 +234,7 @@ import static com.facebook.presto.metadata.MetadataUtil.getConnectorIdOrThrow; import static com.facebook.presto.metadata.MetadataUtil.toSchemaTableName; import static com.facebook.presto.spi.StandardErrorCode.DATATYPE_MISMATCH; -import static com.facebook.presto.spi.StandardErrorCode.INVALID_ARGUMENTS; import static com.facebook.presto.spi.StandardErrorCode.INVALID_COLUMN_MASK; -import static com.facebook.presto.spi.StandardErrorCode.INVALID_FUNCTION_ARGUMENT; import static com.facebook.presto.spi.StandardErrorCode.INVALID_ROW_FILTER; import static com.facebook.presto.spi.StandardWarningCode.PERFORMANCE_WARNING; import static com.facebook.presto.spi.StandardWarningCode.REDUNDANT_ORDER_BY; @@ -243,12 +267,16 @@ import static com.facebook.presto.sql.analyzer.RefreshMaterializedViewPredicateAnalyzer.extractTablePredicates; import static com.facebook.presto.sql.analyzer.ScopeReferenceExtractor.hasReferencesToScope; import static com.facebook.presto.sql.analyzer.SemanticErrorCode.AMBIGUOUS_ATTRIBUTE; +import static com.facebook.presto.sql.analyzer.SemanticErrorCode.AMBIGUOUS_RETURN_TYPE; import static com.facebook.presto.sql.analyzer.SemanticErrorCode.COLUMN_NAME_NOT_SPECIFIED; import static com.facebook.presto.sql.analyzer.SemanticErrorCode.COLUMN_TYPE_UNKNOWN; import static com.facebook.presto.sql.analyzer.SemanticErrorCode.DUPLICATE_COLUMN_NAME; import static com.facebook.presto.sql.analyzer.SemanticErrorCode.DUPLICATE_PARAMETER_NAME; import static com.facebook.presto.sql.analyzer.SemanticErrorCode.DUPLICATE_PROPERTY; import static com.facebook.presto.sql.analyzer.SemanticErrorCode.DUPLICATE_RELATION; +import static com.facebook.presto.sql.analyzer.SemanticErrorCode.FUNCTION_NOT_FOUND; +import static com.facebook.presto.sql.analyzer.SemanticErrorCode.INVALID_ARGUMENTS; +import static com.facebook.presto.sql.analyzer.SemanticErrorCode.INVALID_FUNCTION_ARGUMENT; import static com.facebook.presto.sql.analyzer.SemanticErrorCode.INVALID_FUNCTION_NAME; import static com.facebook.presto.sql.analyzer.SemanticErrorCode.INVALID_OFFSET_ROW_COUNT; import static com.facebook.presto.sql.analyzer.SemanticErrorCode.INVALID_ORDINAL; @@ -258,15 +286,18 @@ import static com.facebook.presto.sql.analyzer.SemanticErrorCode.MATERIALIZED_VIEW_IS_RECURSIVE; import static com.facebook.presto.sql.analyzer.SemanticErrorCode.MISMATCHED_COLUMN_ALIASES; import static com.facebook.presto.sql.analyzer.SemanticErrorCode.MISMATCHED_SET_COLUMN_TYPES; +import static com.facebook.presto.sql.analyzer.SemanticErrorCode.MISSING_ARGUMENT; import static com.facebook.presto.sql.analyzer.SemanticErrorCode.MISSING_ATTRIBUTE; import static com.facebook.presto.sql.analyzer.SemanticErrorCode.MISSING_COLUMN; import static com.facebook.presto.sql.analyzer.SemanticErrorCode.MISSING_MATERIALIZED_VIEW; +import static com.facebook.presto.sql.analyzer.SemanticErrorCode.MISSING_RETURN_TYPE; import static com.facebook.presto.sql.analyzer.SemanticErrorCode.MISSING_SCHEMA; import static com.facebook.presto.sql.analyzer.SemanticErrorCode.MISSING_TABLE; import static com.facebook.presto.sql.analyzer.SemanticErrorCode.MUST_BE_WINDOW_FUNCTION; import static com.facebook.presto.sql.analyzer.SemanticErrorCode.NESTED_WINDOW; import static com.facebook.presto.sql.analyzer.SemanticErrorCode.NONDETERMINISTIC_ORDER_BY_EXPRESSION_WITH_SELECT_DISTINCT; import static com.facebook.presto.sql.analyzer.SemanticErrorCode.NON_NUMERIC_SAMPLE_PERCENTAGE; +import static com.facebook.presto.sql.analyzer.SemanticErrorCode.NOT_IMPLEMENTED; import static com.facebook.presto.sql.analyzer.SemanticErrorCode.NOT_SUPPORTED; import static com.facebook.presto.sql.analyzer.SemanticErrorCode.ORDER_BY_MUST_BE_IN_SELECT; import static com.facebook.presto.sql.analyzer.SemanticErrorCode.TABLE_ALREADY_EXISTS; @@ -465,6 +496,20 @@ protected Scope visitInsert(Insert insert, Optional scope) tableColumnsMetadata.getTableHandle().get(), insertColumns.stream().map(columnHandles::get).collect(toImmutableList()))); + List types = queryScope.getRelationType().getVisibleFields().stream() + .map(Field::getType) + .collect(toImmutableList()); + + Stream columnStream = Streams.zip( + insertColumns.stream(), + types.stream() + .map(Type::toString), + Column::new); + + analysis.setUpdatedSourceColumns(Optional.of(Streams.zip( + columnStream, queryScope.getRelationType().getVisibleFields().stream(), (column, field) -> new OutputColumnMetadata(column.getName(), column.getType(), analysis.getSourceColumns(field))) + .collect(toImmutableList()))); + return createAndAssignScope(insert, scope, Field.newUnqualified(insert.getLocation(), "rows", BIGINT)); } @@ -693,23 +738,36 @@ protected Scope visitCreateTableAsSelect(CreateTableAsSelect node, Optional outputColumns = ImmutableList.builder(); + if (node.getColumnAliases().isPresent()) { validateColumnAliases(node.getColumnAliases().get(), queryScope.getRelationType().getVisibleFieldCount()); - + int aliasPosition = 0; // analyze only column types in subquery if column alias exists for (Field field : queryScope.getRelationType().getVisibleFields()) { if (field.getType().equals(UNKNOWN)) { throw new SemanticException(COLUMN_TYPE_UNKNOWN, node, "Column type is unknown at position %s", queryScope.getRelationType().indexOf(field) + 1); } + String columnName = node.getColumnAliases().get().get(aliasPosition).getValue(); + outputColumns.add(new OutputColumnMetadata(columnName, field.getType().toString(), analysis.getSourceColumns(field))); + aliasPosition++; } } else { validateColumns(node, queryScope.getRelationType()); + queryScope.getRelationType().getVisibleFields().stream() + .map(this::createOutputColumn) + .forEach(outputColumns::add); } - + analysis.setUpdatedSourceColumns(Optional.of(outputColumns.build())); return createAndAssignScope(node, scope, Field.newUnqualified(node.getLocation(), "rows", BIGINT)); } + private OutputColumnMetadata createOutputColumn(Field field) + { + return new OutputColumnMetadata(field.getName().get(), field.getType().toString(), analysis.getSourceColumns(field)); + } + @Override protected Scope visitCreateView(CreateView node, Optional scope) { @@ -1180,7 +1238,7 @@ protected Scope visitExplain(Explain node, Optional scope) .filter(option -> option instanceof ExplainFormat) .map(ExplainFormat.class::cast) .map(ExplainFormat::getType) - .collect(Collectors.toList()); + .collect(toImmutableList()); checkState(formats.size() <= 1, "only a single format option is supported in EXPLAIN ANALYZE"); formats.stream().findFirst().ifPresent(format -> checkState(format.equals(TEXT) || format.equals(JSON), "only TEXT and JSON formats are supported in EXPLAIN ANALYZE")); @@ -1255,7 +1313,7 @@ else if (expressionType instanceof MapType) { outputFields.add(Field.newUnqualified(expression.getLocation(), Optional.empty(), ((MapType) expressionType).getValueType())); } else { - throw new PrestoException(INVALID_FUNCTION_ARGUMENT, "Cannot unnest type: " + expressionType); + throw new PrestoException(StandardErrorCode.INVALID_FUNCTION_ARGUMENT, "Cannot unnest type: " + expressionType); } } if (node.isWithOrdinality()) { @@ -1272,6 +1330,265 @@ protected Scope visitLateral(Lateral node, Optional scope) return createAndAssignScope(node, scope, queryScope.getRelationType()); } + @Override + protected Scope visitTableFunctionInvocation(TableFunctionInvocation node, Optional scope) + { + TableFunctionMetadata tableFunctionMetadata = metadata.getFunctionAndTypeManager() + .getTableFunctionRegistry() + .resolve(session, node.getName()) + .orElseThrow(() -> new SemanticException( + FUNCTION_NOT_FOUND, + node, + "Table function %s not registered", + node.getName())); + + ConnectorTableFunction function = tableFunctionMetadata.getFunction(); + ConnectorId connectorId = tableFunctionMetadata.getConnectorId(); + + QualifiedObjectName functionName = new QualifiedObjectName(connectorId.getCatalogName(), function.getSchema(), function.getName()); + + Map passedArguments = analyzeArguments(node, function.getArguments(), node.getArguments()); + + TransactionManager transactionManager = metadata.getFunctionAndTypeManager().getTransactionManager(); + CatalogMetadata registrationCatalogMetadata = transactionManager.getOptionalCatalogMetadata(session.getRequiredTransactionId(), connectorId.getCatalogName()).orElseThrow(() -> new IllegalStateException("Missing catalog metadata")); + // a call to getRequiredCatalogHandle() is necessary so that the catalog is recorded by the TransactionManager + ConnectorTransactionHandle transactionHandle = transactionManager.getConnectorTransaction( + session.getRequiredTransactionId(), registrationCatalogMetadata.getConnectorId()); + + TableFunctionAnalysis functionAnalysis = function.analyze(session.toConnectorSession(connectorId), transactionHandle, passedArguments); + analysis.setTableFunctionAnalysis(node, new Analysis.TableFunctionInvocationAnalysis(connectorId, functionName.toString(), passedArguments, functionAnalysis.getHandle(), transactionHandle)); + + // TODO process the copartitioning: + // 1. validate input table references + // 2. the copartitioned tables in each set must be partitioned, and have the same number of partitioning columns + // 3. the corresponding columns must be comparable + // 4. within a set, determine and record coercions of the corresponding columns to a common supertype + // Note that if a table is part of multiple copartitioning sets, it might require a different coercion for a column + // per each set. Additionally, there might be another coercion required by the Table Function logic. Also, since + // all partitioning columns are passed-through, we also need an un-coerced copy. + // See ExpressionAnalyzer.sortKeyCoercionsForFrameBoundCalculation for multiple coercions on a column. + if (!node.getCopartitioning().isEmpty()) { + throw new SemanticException(NOT_IMPLEMENTED, node, "COPARTITION clause is not yet supported for table functions"); + } + + // determine the result relation type. + // The result relation type of a table function consists of: + // 1. passed columns from input tables: + // - for tables with the "pass through columns" option, these are all columns of the table, + // - for tables without the "pass through columns" option, these are the partitioning columns of the table, if any. + // 2. columns created by the table function, called the proper columns. + ReturnTypeSpecification returnTypeSpecification = function.getReturnTypeSpecification(); + Optional analyzedProperColumnsDescriptor = functionAnalysis.getReturnedType(); + Descriptor properColumnsDescriptor; + switch (returnTypeSpecification.getReturnType()) { + case ReturnTypeSpecification.OnlyPassThrough.returnType: + throw new SemanticException(NOT_IMPLEMENTED, node, "Returning only pass through columns is not yet supported for table functions"); + case ReturnTypeSpecification.GenericTable.returnType: + properColumnsDescriptor = analyzedProperColumnsDescriptor + .orElseThrow(() -> new SemanticException(MISSING_RETURN_TYPE, node, "Cannot determine returned relation type for table function " + node.getName())); + break; + default: + // returned type is statically declared at function declaration and cannot be overridden + if (analyzedProperColumnsDescriptor.isPresent()) { + throw new SemanticException(AMBIGUOUS_RETURN_TYPE, node, "Returned relation type for table function %s is ambiguous", node.getName()); + } + properColumnsDescriptor = ((ReturnTypeSpecification.DescribedTable) returnTypeSpecification).getDescriptor(); + } + // currently we don't support input tables, so the output consists of proper columns only + List fields = properColumnsDescriptor.getFields().stream() + // per spec, field names are mandatory + .map(field -> Field.newUnqualified((field.getName()), field.getType().orElseThrow(() -> new IllegalStateException("missing returned type for proper field")))) + .collect(toImmutableList()); + + return createAndAssignScope(node, scope, fields); + } + + private Map analyzeArguments(Node node, List argumentSpecifications, List arguments) + { + Node errorLocation = node; + if (!arguments.isEmpty()) { + errorLocation = arguments.get(0); + } + + if (argumentSpecifications.size() < arguments.size()) { + throw new SemanticException(INVALID_ARGUMENTS, errorLocation, "Too many arguments. Expected at most %s arguments, got %s arguments", argumentSpecifications.size(), arguments.size()); + } + + if (argumentSpecifications.isEmpty()) { + return ImmutableMap.of(); + } + + boolean argumentsPassedByName = !arguments.isEmpty() && arguments.stream().allMatch(argument -> argument.getName().isPresent()); + boolean argumentsPassedByPosition = arguments.stream().allMatch(argument -> !argument.getName().isPresent()); + if (!argumentsPassedByName && !argumentsPassedByPosition) { + throw new SemanticException(INVALID_ARGUMENTS, errorLocation, "All arguments must be passed by name or all must be passed positionally"); + } + + if (argumentsPassedByName) { + return mapTableFunctionsArgsByName(argumentSpecifications, arguments, errorLocation); + } + else { + return mapTableFunctionArgsByPosition(argumentSpecifications, arguments, errorLocation); + } + } + + private Map mapTableFunctionsArgsByName(List argumentSpecifications, List arguments, Node errorLocation) + { + ImmutableMap.Builder passedArguments = ImmutableMap.builder(); + Map argumentSpecificationsByName = new HashMap<>(); + for (ArgumentSpecification argumentSpecification : argumentSpecifications) { + if (argumentSpecificationsByName.put(argumentSpecification.getName(), argumentSpecification) != null) { + // this should never happen, because the argument names are validated at function registration time + throw new IllegalStateException("Duplicate argument specification for name: " + argumentSpecification.getName()); + } + } + Set uniqueArgumentNames = new HashSet<>(); + for (TableFunctionArgument argument : arguments) { + String argumentName = argument.getName().orElseThrow(() -> new IllegalStateException("Missing table function argument name")).getCanonicalValue(); + if (!uniqueArgumentNames.add(argumentName)) { + throw new SemanticException(INVALID_FUNCTION_ARGUMENT, argument, "Duplicate argument name: %s", argumentName); + } + ArgumentSpecification argumentSpecification = argumentSpecificationsByName.remove(argumentName); + if (argumentSpecification == null) { + throw new SemanticException(INVALID_FUNCTION_ARGUMENT, argument, "Unexpected argument name: %s", argumentName); + } + passedArguments.put(argumentSpecification.getName(), analyzeArgument(argumentSpecification, argument)); + } + // apply defaults for not specified arguments + for (Map.Entry entry : argumentSpecificationsByName.entrySet()) { + ArgumentSpecification argumentSpecification = entry.getValue(); + passedArguments.put(argumentSpecification.getName(), analyzeDefault(argumentSpecification, errorLocation)); + } + return passedArguments.build(); + } + + private Map mapTableFunctionArgsByPosition(List argumentSpecifications, List arguments, Node errorLocation) + { + ImmutableMap.Builder passedArguments = ImmutableMap.builder(); + for (int i = 0; i < arguments.size(); i++) { + TableFunctionArgument argument = arguments.get(i); + ArgumentSpecification argumentSpecification = argumentSpecifications.get(i); // TODO args passed positionally - can one only pass some prefix of args? + passedArguments.put(argumentSpecification.getName(), analyzeArgument(argumentSpecification, argument)); + } + // apply defaults for not specified arguments + for (int i = arguments.size(); i < argumentSpecifications.size(); i++) { + ArgumentSpecification argumentSpecification = argumentSpecifications.get(i); + passedArguments.put(argumentSpecification.getName(), analyzeDefault(argumentSpecification, errorLocation)); + } + return passedArguments.build(); + } + + private Argument analyzeArgument(ArgumentSpecification argumentSpecification, TableFunctionArgument argument) + { + String actualType = getArgumentTypeString(argument); + switch (argumentSpecification.getArgumentType()){ + case TableArgumentSpecification.argumentType: + return analyzeTableArgument(argument, argumentSpecification, actualType); + case DescriptorArgumentSpecification.argumentType: + return analyzeDescriptorArgument(argument, argumentSpecification, actualType); + case ScalarArgumentSpecification.argumentType: + return analyzeScalarArgument(argument, argumentSpecification, actualType); + default: + throw new IllegalStateException("Unexpected argument specification: " + argumentSpecification.getClass().getSimpleName()); + } + } + + private Argument analyzeDefault(ArgumentSpecification argumentSpecification, Node errorLocation) + { + if (argumentSpecification.isRequired()) { + throw new SemanticException(MISSING_ARGUMENT, errorLocation, "Missing argument: " + argumentSpecification.getName()); + } + + checkArgument(!(argumentSpecification instanceof TableArgumentSpecification), "invalid table argument specification: default set"); + + if (argumentSpecification instanceof DescriptorArgumentSpecification) { + throw new SemanticException(NOT_IMPLEMENTED, errorLocation, "Descriptor arguments are not yet supported for table functions"); + } + if (argumentSpecification instanceof ScalarArgumentSpecification) { + return ScalarArgument.builder() + .type(((ScalarArgumentSpecification) argumentSpecification).getType()) + .value(argumentSpecification.getDefaultValue()) + .build(); + } + + throw new IllegalStateException("Unexpected argument specification: " + argumentSpecification.getClass().getSimpleName()); + } + + private String getArgumentTypeString(TableFunctionArgument argument) + { + try { + return argument.getValue().getArgumentTypeString(); + } + catch (IllegalArgumentException e) { + throw new SemanticException(INVALID_FUNCTION_ARGUMENT, argument, "Unexpected table function argument type: ", argument.getClass().getSimpleName()); + } + } + + private Argument analyzeScalarArgument(TableFunctionArgument argument, ArgumentSpecification argumentSpecification, String actualType) + { + Type type = ((ScalarArgumentSpecification) argumentSpecification).getType(); + if (!(argument.getValue() instanceof Expression)) { + throw new SemanticException(INVALID_FUNCTION_ARGUMENT, argument, "Invalid argument %s. Expected expression, got %s", argumentSpecification.getName(), actualType); + } + Expression expression = (Expression) argument.getValue(); + // 'descriptor' as a function name is not allowed in this context + if (expression instanceof FunctionCall && ((FunctionCall) expression).getName().hasSuffix(QualifiedName.of("descriptor"))) { // function name is always compared case-insensitive + throw new SemanticException(INVALID_FUNCTION_ARGUMENT, argument, "'descriptor' function is not allowed as a table function argument"); + } + // inline parameters + Expression inlined = ExpressionTreeRewriter.rewriteWith(new ExpressionRewriter() + { + @Override + public Expression rewriteParameter(Parameter node, Void context, ExpressionTreeRewriter treeRewriter) + { + if (analysis.isDescribe()) { + // We cannot handle DESCRIBE when a table function argument involves a parameter. + // In DESCRIBE, the parameter values are not known. We cannot pass a dummy value for a parameter. + // The value of a table function argument can affect the returned relation type. The returned + // relation type can affect the assumed types for other parameters in the query. + throw new SemanticException(NOT_SUPPORTED, node, "DESCRIBE is not supported if a table function uses parameters"); + } + return analysis.getParameters().get(NodeRef.of(node)); + } + }, expression); + // currently, only constant arguments are supported + Object constantValue = ExpressionInterpreter.evaluateConstantExpression(inlined, type, metadata, session, analysis.getParameters()); + return ScalarArgument.builder() + .type(type) + .value(constantValue) + .build(); + } + + private Argument analyzeTableArgument(TableFunctionArgument argument, ArgumentSpecification argumentSpecification, String actualType) + { + if (!(argument.getValue() instanceof TableFunctionTableArgument)) { + if (argument.getValue() instanceof FunctionCall) { + // probably an attempt to pass a table function call, which is not supported, and was parsed as a function call + throw new SemanticException(NOT_IMPLEMENTED, argument, "Invalid table argument %s. Table functions are not allowed as table function arguments", argumentSpecification.getName()); + } + throw new SemanticException(INVALID_FUNCTION_ARGUMENT, argument, "Invalid argument %s. Expected table, got %s", argumentSpecification.getName(), actualType); + } + // TODO analyze the argument + // 1. process the Relation + // 2. partitioning and ordering must only apply to tables with set semantics + // 3. validate partitioning and ordering using `validateAndGetInputField()` + // 4. validate the prune when empty property vs argument specification (forbidden for row semantics; override? -> check spec) + // 5. return Argument + throw new SemanticException(NOT_IMPLEMENTED, argument, "Table arguments are not yet supported for table functions"); + } + + private Argument analyzeDescriptorArgument(TableFunctionArgument argument, ArgumentSpecification argumentSpecification, String actualType) + { + if (!(argument.getValue() instanceof TableFunctionDescriptorArgument)) { + if (argument.getValue() instanceof FunctionCall && ((FunctionCall) argument.getValue()).getName().hasSuffix(QualifiedName.of("descriptor"))) { // function name is always compared case-insensitive + // malformed descriptor which parsed as a function call + throw new SemanticException(INVALID_FUNCTION_ARGUMENT, argument, "Invalid descriptor argument %s. Descriptors should be formatted as 'DESCRIPTOR(name [type], ...)'", (Object) argumentSpecification.getName()); + } + throw new SemanticException(INVALID_FUNCTION_ARGUMENT, argument, "Invalid argument %s. Expected descriptor, got %s", argumentSpecification.getName(), actualType); + } + throw new SemanticException(NOT_IMPLEMENTED, argument, "Descriptor arguments are not yet supported for table functions"); + } + @Override protected Scope visitTable(Table table, Optional scope) { @@ -1296,7 +1613,7 @@ protected Scope visitTable(Table table, Optional scope) Iterator visibleFieldsIterator = queryDescriptor.getVisibleFields().iterator(); for (Identifier columnName : columnNames.get()) { Field inputField = visibleFieldsIterator.next(); - fieldBuilder.add(Field.newQualified( + Field field = Field.newQualified( columnName.getLocation(), QualifiedName.of(name), Optional.of(columnName.getValue()), @@ -1304,23 +1621,29 @@ protected Scope visitTable(Table table, Optional scope) false, inputField.getOriginTable(), inputField.getOriginColumnName(), - inputField.isAliased())); + inputField.isAliased()); + fieldBuilder.add(field); + analysis.addSourceColumns(field, analysis.getSourceColumns(inputField)); } fields = fieldBuilder.build(); } else { - fields = queryDescriptor.getAllFields().stream() - .map(field -> Field.newQualified( - field.getNodeLocation(), - QualifiedName.of(name), - field.getName(), - field.getType(), - field.isHidden(), - field.getOriginTable(), - field.getOriginColumnName(), - field.isAliased())) - .collect(toImmutableList()); + ImmutableList.Builder fieldBuilder = ImmutableList.builder(); + for (Field inputField : queryDescriptor.getAllFields()) { + Field field = Field.newQualified( + inputField.getNodeLocation(), + QualifiedName.of(name), + inputField.getName(), + inputField.getType(), + inputField.isHidden(), + inputField.getOriginTable(), + inputField.getOriginColumnName(), + inputField.isAliased()); + fieldBuilder.add(field); + analysis.addSourceColumns(field, analysis.getSourceColumns(inputField)); + } + fields = fieldBuilder.build(); } return createAndAssignScope(table, scope, fields); @@ -1388,6 +1711,7 @@ protected Scope visitTable(Table table, Optional scope) ColumnHandle columnHandle = columnHandles.get(column.getName()); checkArgument(columnHandle != null, "Unknown field %s", field); analysis.setColumn(field, columnHandle); + analysis.addSourceColumns(field, ImmutableSet.of(new SourceColumn(name, column.getName()))); } analysis.registerTable(table, tableHandle.get()); @@ -1458,7 +1782,7 @@ private Optional processTableVersion(Table table, QualifiedObjectNa analysis.recordSubqueries(table, expressionAnalysis); Type stateExprType = expressionAnalysis.getType(stateExpr); if (stateExprType == UNKNOWN) { - throw new PrestoException(INVALID_ARGUMENTS, format("Table version AS OF/BEFORE expression cannot be NULL for %s", name.toString())); + throw new PrestoException(StandardErrorCode.INVALID_ARGUMENTS, format("Table version AS OF/BEFORE expression cannot be NULL for %s", name.toString())); } Object evalStateExpr = evaluateConstantExpression(stateExpr, stateExprType, metadata, session, analysis.getParameters()); if (tableVersionType == TIMESTAMP) { @@ -1712,13 +2036,24 @@ protected Scope visitAliasedRelation(AliasedRelation relation, Optional s } List aliases = null; + Collection inputFields = relationType.getAllFields(); if (relation.getColumnNames() != null) { aliases = relation.getColumnNames().stream() .map(Identifier::getValue) - .collect(Collectors.toList()); + .collect(toImmutableList()); + inputFields = relationType.getVisibleFields(); } RelationType descriptor = relationType.withAlias(relation.getAlias().getValue(), aliases); + checkArgument(inputFields.size() == descriptor.getAllFieldCount(), + "Expected %s fields, got %s", + descriptor.getAllFieldCount(), + inputFields.size()); + + Streams.forEachPair( + descriptor.getAllFields().stream(), + inputFields.stream(), + (newField, field) -> analysis.addSourceColumns(newField, analysis.getSourceColumns(field))); return createAndAssignScope(relation, scope, descriptor); } @@ -1918,6 +2253,14 @@ protected Scope visitSetOperation(SetOperation node, Optional scope) oldField.getOriginTable(), oldField.getOriginColumnName(), oldField.isAliased()); + + int index = i; + analysis.addSourceColumns( + outputDescriptorFields[index], + relationScopes.stream() + .map(relationType -> relationType.getRelationType().getFieldByIndex(index)) + .flatMap(field -> analysis.getSourceColumns(field).stream()) + .collect(toImmutableSet())); } for (int i = 0; i < node.getRelations().size(); i++) { @@ -2478,7 +2821,7 @@ private void checkFunctionName(Statement node, QualifiedName functionName, boole .map(SqlFunction::getSignature) .map(Signature::getName) .map(QualifiedObjectName::getObjectName) - .collect(Collectors.toList()); + .collect(toImmutableList()); if (builtInFunctionNames.contains(functionName.toString())) { throw new SemanticException(INVALID_FUNCTION_NAME, node, format("Function %s is already registered as a built-in function.", functionName)); } @@ -2668,7 +3011,9 @@ private Scope computeAndAssignOutputScope(QuerySpecification node, Optional starPrefix = ((AllColumns) item).getPrefix(); for (Field field : sourceScope.getRelationType().resolveFieldsWithPrefix(starPrefix)) { - outputFields.add(Field.newUnqualified(node.getSelect().getLocation(), field.getName(), field.getType(), field.getOriginTable(), field.getOriginColumnName(), false)); + Field newField = Field.newUnqualified(node.getSelect().getLocation(), field.getName(), field.getType(), field.getOriginTable(), field.getOriginColumnName(), false); + analysis.addSourceColumns(newField, analysis.getSourceColumns(field)); + outputFields.add(newField); } } else if (item instanceof SingleColumn) { @@ -2701,8 +3046,14 @@ else if (expression instanceof DereferenceExpression) { field = Optional.of(name.getOriginalSuffix()); } } - - outputFields.add(Field.newUnqualified(expression.getLocation(), field.map(Identifier::getValue), analysis.getType(expression), originTable, originColumn, column.getAlias().isPresent())); // TODO don't use analysis as a side-channel. Use outputExpressions to look up the type + Field newField = Field.newUnqualified(expression.getLocation(), field.map(Identifier::getValue), analysis.getType(expression), originTable, originColumn, column.getAlias().isPresent()); + if (originTable.isPresent()) { + analysis.addSourceColumns(newField, ImmutableSet.of(new SourceColumn(originTable.get(), originColumn.orElseThrow()))); + } + else { + analysis.addSourceColumns(newField, analysis.getExpressionSourceColumns(expression)); + } + outputFields.add(newField); } else { throw new IllegalArgumentException("Unsupported SelectItem type: " + item.getClass().getName()); diff --git a/presto-main-base/src/main/java/com/facebook/presto/sql/planner/CanonicalPlanGenerator.java b/presto-main-base/src/main/java/com/facebook/presto/sql/planner/CanonicalPlanGenerator.java index 05fc2df18692d..6f5274eb36925 100644 --- a/presto-main-base/src/main/java/com/facebook/presto/sql/planner/CanonicalPlanGenerator.java +++ b/presto-main-base/src/main/java/com/facebook/presto/sql/planner/CanonicalPlanGenerator.java @@ -20,6 +20,7 @@ import com.facebook.presto.spi.PrestoException; import com.facebook.presto.spi.SchemaTableName; import com.facebook.presto.spi.VariableAllocator; +import com.facebook.presto.spi.eventlistener.OutputColumnMetadata; import com.facebook.presto.spi.plan.AggregationNode; import com.facebook.presto.spi.plan.AggregationNode.Aggregation; import com.facebook.presto.spi.plan.AggregationNode.GroupingSetDescriptor; @@ -1115,6 +1116,12 @@ public SchemaTableName getSchemaTableName() return new SchemaTableName("schema", "table"); } + @Override + public Optional> getOutputColumns() + { + return Optional.empty(); + } + @Override public String toString() { diff --git a/presto-main-base/src/main/java/com/facebook/presto/sql/planner/LocalExecutionPlanner.java b/presto-main-base/src/main/java/com/facebook/presto/sql/planner/LocalExecutionPlanner.java index c07fd31994d6b..85a52a298bcb7 100644 --- a/presto-main-base/src/main/java/com/facebook/presto/sql/planner/LocalExecutionPlanner.java +++ b/presto-main-base/src/main/java/com/facebook/presto/sql/planner/LocalExecutionPlanner.java @@ -210,6 +210,7 @@ import com.facebook.presto.sql.planner.plan.RowNumberNode; import com.facebook.presto.sql.planner.plan.SampleNode; import com.facebook.presto.sql.planner.plan.StatisticsWriterNode; +import com.facebook.presto.sql.planner.plan.TableFunctionNode; import com.facebook.presto.sql.planner.plan.TableWriterMergeNode; import com.facebook.presto.sql.planner.plan.TopNRowNumberNode; import com.facebook.presto.sql.planner.plan.UpdateNode; @@ -1204,6 +1205,12 @@ public PhysicalOperation visitWindow(WindowNode node, LocalExecutionPlanContext return new PhysicalOperation(operatorFactory, outputMappings.build(), context, source); } + @Override + public PhysicalOperation visitTableFunction(TableFunctionNode node, LocalExecutionPlanContext context) + { + throw new UnsupportedOperationException("execution by operator is not yet implemented for table function " + node.getName()); + } + @Override public PhysicalOperation visitTopN(TopNNode node, LocalExecutionPlanContext context) { diff --git a/presto-main-base/src/main/java/com/facebook/presto/sql/planner/LogicalPlanner.java b/presto-main-base/src/main/java/com/facebook/presto/sql/planner/LogicalPlanner.java index cf748ba0eb44e..e1c1979f924be 100644 --- a/presto-main-base/src/main/java/com/facebook/presto/sql/planner/LogicalPlanner.java +++ b/presto-main-base/src/main/java/com/facebook/presto/sql/planner/LogicalPlanner.java @@ -286,7 +286,7 @@ private RelationPlan createTableCreationPlan(Analysis analysis, Query query) return createTableWriterPlan( analysis, plan, - new CreateName(new ConnectorId(destination.getCatalogName()), tableMetadata, newTableLayout), + new CreateName(new ConnectorId(destination.getCatalogName()), tableMetadata, newTableLayout, analysis.getUpdatedSourceColumns()), columnNames, tableMetadata.getColumns(), newTableLayout, @@ -310,7 +310,7 @@ private RelationPlan createInsertPlan(Analysis analysis, Insert insertStatement) TableHandle tableHandle = insertAnalysis.getTarget(); List columnHandles = insertAnalysis.getColumns(); - TableWriterNode.WriterTarget target = new InsertReference(tableHandle, metadata.getTableMetadata(session, tableHandle).getTable()); + TableWriterNode.WriterTarget target = new InsertReference(tableHandle, metadata.getTableMetadata(session, tableHandle).getTable(), analysis.getUpdatedSourceColumns()); return buildInternalInsertPlan(tableHandle, columnHandles, insertStatement.getQuery(), analysis, target); } diff --git a/presto-main-base/src/main/java/com/facebook/presto/sql/planner/OutputExtractor.java b/presto-main-base/src/main/java/com/facebook/presto/sql/planner/OutputExtractor.java index bd4418d337e1f..97c95e0d7bf05 100644 --- a/presto-main-base/src/main/java/com/facebook/presto/sql/planner/OutputExtractor.java +++ b/presto-main-base/src/main/java/com/facebook/presto/sql/planner/OutputExtractor.java @@ -13,23 +13,20 @@ */ package com.facebook.presto.sql.planner; -import com.facebook.presto.execution.Column; import com.facebook.presto.execution.Output; import com.facebook.presto.spi.ConnectorId; import com.facebook.presto.spi.SchemaTableName; +import com.facebook.presto.spi.eventlistener.OutputColumnMetadata; import com.facebook.presto.spi.plan.PlanNode; import com.facebook.presto.spi.plan.TableWriterNode; import com.facebook.presto.sql.planner.plan.InternalPlanVisitor; import com.facebook.presto.sql.planner.plan.SequenceNode; import com.google.common.base.VerifyException; -import com.google.common.collect.ImmutableList; -import java.util.ArrayList; import java.util.List; import java.util.Optional; import static com.facebook.presto.spi.connector.ConnectorCommitHandle.EMPTY_COMMIT_OUTPUT; -import static com.google.common.base.Preconditions.checkArgument; import static com.google.common.base.Preconditions.checkState; public class OutputExtractor @@ -47,7 +44,7 @@ public Optional extractOutput(PlanNode root) visitor.getSchemaTableName().getSchemaName(), visitor.getSchemaTableName().getTableName(), EMPTY_COMMIT_OUTPUT, - Optional.of(ImmutableList.copyOf(visitor.getColumns())))); + visitor.getOutputColumns())); } private class Visitor @@ -55,7 +52,7 @@ private class Visitor { private ConnectorId connectorId; private SchemaTableName schemaTableName; - private List columns = new ArrayList<>(); + private Optional> outputColumns = Optional.empty(); @Override public Void visitTableWriter(TableWriterNode node, Void context) @@ -65,11 +62,7 @@ public Void visitTableWriter(TableWriterNode node, Void context) checkState(schemaTableName == null || schemaTableName.equals(writerTarget.getSchemaTableName()), "cannot have more than a single create, insert or delete in a query"); schemaTableName = writerTarget.getSchemaTableName(); - - checkArgument(node.getColumnNames().size() == node.getColumns().size(), "Column names and columns sizes must be equal"); - for (int i = 0; i < node.getColumnNames().size(); i++) { - columns.add(new Column(node.getColumnNames().get(i), node.getColumns().get(i).getType().toString())); - } + outputColumns = writerTarget.getOutputColumns(); return null; } @@ -100,9 +93,9 @@ public SchemaTableName getSchemaTableName() return schemaTableName; } - public List getColumns() + public Optional> getOutputColumns() { - return columns; + return outputColumns; } } } diff --git a/presto-main-base/src/main/java/com/facebook/presto/sql/planner/PlanOptimizers.java b/presto-main-base/src/main/java/com/facebook/presto/sql/planner/PlanOptimizers.java index 4d08e779911c9..927fda540e07a 100644 --- a/presto-main-base/src/main/java/com/facebook/presto/sql/planner/PlanOptimizers.java +++ b/presto-main-base/src/main/java/com/facebook/presto/sql/planner/PlanOptimizers.java @@ -133,6 +133,7 @@ import com.facebook.presto.sql.planner.iterative.rule.RewriteConstantArrayContainsToInExpression; import com.facebook.presto.sql.planner.iterative.rule.RewriteFilterWithExternalFunctionToProject; import com.facebook.presto.sql.planner.iterative.rule.RewriteSpatialPartitioningAggregation; +import com.facebook.presto.sql.planner.iterative.rule.RewriteTableFunctionToTableScan; import com.facebook.presto.sql.planner.iterative.rule.RuntimeReorderJoinSides; import com.facebook.presto.sql.planner.iterative.rule.ScaledWriterRule; import com.facebook.presto.sql.planner.iterative.rule.SimplifyCardinalityMap; @@ -855,6 +856,14 @@ public PlanOptimizers( costCalculator, ImmutableSet.of(new ScaledWriterRule()))); + builder.add( + new IterativeOptimizer( + metadata, + ruleStats, + statsCalculator, + costCalculator, + ImmutableSet.of(new RewriteTableFunctionToTableScan(metadata)))); + if (!noExchange) { builder.add(new ReplicateSemiJoinInDelete()); // Must run before AddExchanges diff --git a/presto-main-base/src/main/java/com/facebook/presto/sql/planner/RelationPlanner.java b/presto-main-base/src/main/java/com/facebook/presto/sql/planner/RelationPlanner.java index 81e6e0bd1598d..11241f2f3dce1 100644 --- a/presto-main-base/src/main/java/com/facebook/presto/sql/planner/RelationPlanner.java +++ b/presto-main-base/src/main/java/com/facebook/presto/sql/planner/RelationPlanner.java @@ -21,6 +21,7 @@ import com.facebook.presto.common.type.RowType; import com.facebook.presto.common.type.Type; import com.facebook.presto.metadata.Metadata; +import com.facebook.presto.metadata.TableFunctionHandle; import com.facebook.presto.spi.ColumnHandle; import com.facebook.presto.spi.PrestoException; import com.facebook.presto.spi.TableHandle; @@ -55,6 +56,7 @@ import com.facebook.presto.sql.planner.optimizations.SampleNodeUtil; import com.facebook.presto.sql.planner.plan.LateralJoinNode; import com.facebook.presto.sql.planner.plan.SampleNode; +import com.facebook.presto.sql.planner.plan.TableFunctionNode; import com.facebook.presto.sql.tree.AliasedRelation; import com.facebook.presto.sql.tree.Cast; import com.facebook.presto.sql.tree.CoalesceExpression; @@ -84,6 +86,7 @@ import com.facebook.presto.sql.tree.SetOperation; import com.facebook.presto.sql.tree.SymbolReference; import com.facebook.presto.sql.tree.Table; +import com.facebook.presto.sql.tree.TableFunctionInvocation; import com.facebook.presto.sql.tree.TableSubquery; import com.facebook.presto.sql.tree.Union; import com.facebook.presto.sql.tree.Unnest; @@ -297,6 +300,44 @@ private RelationPlan addColumnMasks(Table table, RelationPlan plan, SqlPlannerCo return new RelationPlan(planBuilder.getRoot(), plan.getScope(), newMappings.build()); } + @Override + protected RelationPlan visitTableFunctionInvocation(TableFunctionInvocation node, SqlPlannerContext context) + { + Analysis.TableFunctionInvocationAnalysis functionAnalysis = analysis.getTableFunctionAnalysis(node); + + // TODO handle input relations: + // 1. extract the input relations from node.getArguments() and plan them. Apply relation coercions if requested. + // 2. for each input relation, prepare the TableArgumentProperties record, consisting of: + // - row or set semantics (from the actualArgument) + // - prune when empty property (from the actualArgument) + // - pass through columns property (from the actualArgument) + // - optional Specification: ordering scheme and partitioning (from the node's argument) <- planned upon the source's RelationPlan (or combined RelationPlan from all sources) + // TODO add - argument name + // TODO add - mapping column name => Symbol // TODO mind the fields without names and duplicate field names in RelationType + List sources = ImmutableList.of(); + List inputRelationsProperties = ImmutableList.of(); + + Scope scope = analysis.getScope(node); + + ImmutableList.Builder outputVariablesBuilder = ImmutableList.builder(); + for (Field field : scope.getRelationType().getAllFields()) { + VariableReferenceExpression variable = variableAllocator.newVariable(getSourceLocation(node), field.getName().get(), field.getType()); + outputVariablesBuilder.add(variable); + } + + List outputVariables = outputVariablesBuilder.build(); + PlanNode root = new TableFunctionNode( + idAllocator.getNextId(), + functionAnalysis.getFunctionName(), + functionAnalysis.getArguments(), + outputVariablesBuilder.build(), + sources.stream().map(RelationPlan::getRoot).collect(toImmutableList()), + inputRelationsProperties, + new TableFunctionHandle(functionAnalysis.getConnectorId(), functionAnalysis.getConnectorTableFunctionHandle(), functionAnalysis.getTransactionHandle())); + + return new RelationPlan(root, scope, outputVariables); + } + @Override protected RelationPlan visitAliasedRelation(AliasedRelation node, SqlPlannerContext context) { diff --git a/presto-main-base/src/main/java/com/facebook/presto/sql/planner/iterative/Lookup.java b/presto-main-base/src/main/java/com/facebook/presto/sql/planner/iterative/Lookup.java index 412e35666ddc8..42c860d36b2d8 100644 --- a/presto-main-base/src/main/java/com/facebook/presto/sql/planner/iterative/Lookup.java +++ b/presto-main-base/src/main/java/com/facebook/presto/sql/planner/iterative/Lookup.java @@ -26,7 +26,7 @@ public interface Lookup /** * Resolves a node by materializing GroupReference nodes * representing symbolic references to other nodes. This method - * is deprecated since is assumes group contains only one node. + * is deprecated since it assumes group contains only one node. *

* If the node is not a GroupReference, it returns the * argument as is. diff --git a/presto-main-base/src/main/java/com/facebook/presto/sql/planner/iterative/rule/ExtractSpatialJoins.java b/presto-main-base/src/main/java/com/facebook/presto/sql/planner/iterative/rule/ExtractSpatialJoins.java index 851e40b05825f..4b30f9de7f9aa 100644 --- a/presto-main-base/src/main/java/com/facebook/presto/sql/planner/iterative/rule/ExtractSpatialJoins.java +++ b/presto-main-base/src/main/java/com/facebook/presto/sql/planner/iterative/rule/ExtractSpatialJoins.java @@ -469,6 +469,11 @@ else if (alignment < 0) { private static boolean isSphericalJoin(Metadata metadata, RowExpression firstArgument, RowExpression secondArgument) { + // In sidecar-enabled clusters, SphericalGeography isn't a supported type. + // If SphericalGeography is not supported, it can be assumed that this join isn't a spherical join, hence returning False. + if (!metadata.getFunctionAndTypeManager().hasType(SPHERICAL_GEOGRAPHY_TYPE_SIGNATURE)) { + return false; + } Type sphericalGeographyType = metadata.getType(SPHERICAL_GEOGRAPHY_TYPE_SIGNATURE); return firstArgument.getType().equals(sphericalGeographyType) || secondArgument.getType().equals(sphericalGeographyType); } diff --git a/presto-main-base/src/main/java/com/facebook/presto/sql/planner/iterative/rule/LeftJoinWithArrayContainsToEquiJoinCondition.java b/presto-main-base/src/main/java/com/facebook/presto/sql/planner/iterative/rule/LeftJoinWithArrayContainsToEquiJoinCondition.java index b5348879355cf..bd1e575f8e78b 100644 --- a/presto-main-base/src/main/java/com/facebook/presto/sql/planner/iterative/rule/LeftJoinWithArrayContainsToEquiJoinCondition.java +++ b/presto-main-base/src/main/java/com/facebook/presto/sql/planner/iterative/rule/LeftJoinWithArrayContainsToEquiJoinCondition.java @@ -48,7 +48,7 @@ import static java.util.Objects.requireNonNull; /** - * When the join condition of a left join has pattern of contains(array, element) where array, we can rewrite it as a equi join condition. For example: + * When the join condition of a left join has pattern of contains(array, element) where array is from the right-side relation and element is from the left-side relation, we can rewrite it as an equi join condition. For example: * 

  * - Left Join
  *      empty join clause
diff --git a/presto-main-base/src/main/java/com/facebook/presto/sql/planner/iterative/rule/RewriteTableFunctionToTableScan.java b/presto-main-base/src/main/java/com/facebook/presto/sql/planner/iterative/rule/RewriteTableFunctionToTableScan.java
new file mode 100644
index 0000000000000..2418377c7ac53
--- /dev/null
+++ b/presto-main-base/src/main/java/com/facebook/presto/sql/planner/iterative/rule/RewriteTableFunctionToTableScan.java
@@ -0,0 +1,97 @@
+/*
+ * Licensed 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 com.facebook.presto.sql.planner.iterative.rule;
+
+import com.facebook.presto.common.predicate.TupleDomain;
+import com.facebook.presto.matching.Captures;
+import com.facebook.presto.matching.Pattern;
+import com.facebook.presto.metadata.Metadata;
+import com.facebook.presto.spi.ColumnHandle;
+import com.facebook.presto.spi.TableHandle;
+import com.facebook.presto.spi.connector.TableFunctionApplicationResult;
+import com.facebook.presto.spi.plan.TableScanNode;
+import com.facebook.presto.spi.relation.VariableReferenceExpression;
+import com.facebook.presto.sql.planner.iterative.Rule;
+import com.facebook.presto.sql.planner.plan.TableFunctionNode;
+import com.google.common.collect.ImmutableMap;
+
+import java.util.List;
+import java.util.Optional;
+
+import static com.facebook.presto.matching.Pattern.empty;
+import static com.facebook.presto.sql.planner.plan.Patterns.sources;
+import static com.facebook.presto.sql.planner.plan.Patterns.tableFunction;
+import static com.google.common.base.Preconditions.checkState;
+import static java.util.Objects.requireNonNull;
+
+/*
+ * This process converts connector-resolvable TableFunctionNodes into equivalent
+ * TableScanNodes by invoking the connector’s applyTableFunction() during planning.
+ * It allows table-valued functions whose results can be expressed as a ConnectorTableHandle
+ * to be treated like regular scans and benefit from normal scan optimizations.
+ *
+ * Example:
+ * Before Transformation:
+ *   TableFunction(my_function(arg1, arg2))
+ *
+ * After Transformation:
+ *   TableScan(my_function(arg1, arg2)).applyTableFunction_tableHandle)
+ *    assignments: {outputVar1 -> my_function(arg1, arg2)).applyTableFunction_colHandle1,
+ *                  outputVar2 -> my_function(arg1, arg2)).applyTableFunction_colHandle2}
+ */
+public class RewriteTableFunctionToTableScan
+        implements Rule
+{
+    private static final Pattern PATTERN = tableFunction()
+            .with(empty(sources()));
+
+    private final Metadata metadata;
+
+    public RewriteTableFunctionToTableScan(Metadata metadata)
+    {
+        this.metadata = requireNonNull(metadata, "metadata is null");
+    }
+
+    @Override
+    public Pattern getPattern()
+    {
+        return PATTERN;
+    }
+
+    @Override
+    public Result apply(TableFunctionNode tableFunctionNode, Captures captures, Context context)
+    {
+        Optional> result = metadata.applyTableFunction(context.getSession(), tableFunctionNode.getHandle());
+
+        if (!result.isPresent()) {
+            return Result.empty();
+        }
+
+        List columnHandles = result.get().getColumnHandles();
+        checkState(tableFunctionNode.getOutputVariables().size() == columnHandles.size(), "returned table does not match the node's output");
+        ImmutableMap.Builder assignments = ImmutableMap.builder();
+        for (int i = 0; i < columnHandles.size(); i++) {
+            assignments.put(tableFunctionNode.getOutputVariables().get(i), columnHandles.get(i));
+        }
+
+        return Result.ofPlanNode(new TableScanNode(
+                tableFunctionNode.getSourceLocation(),
+                tableFunctionNode.getId(),
+                result.get().getTableHandle(),
+                tableFunctionNode.getOutputVariables(),
+                assignments.buildOrThrow(),
+                TupleDomain.all(),
+                TupleDomain.all(), Optional.empty()));
+    }
+}
diff --git a/presto-main-base/src/main/java/com/facebook/presto/sql/planner/optimizations/ActualProperties.java b/presto-main-base/src/main/java/com/facebook/presto/sql/planner/optimizations/ActualProperties.java
index 32fc8913fa12c..f562d79b50c32 100644
--- a/presto-main-base/src/main/java/com/facebook/presto/sql/planner/optimizations/ActualProperties.java
+++ b/presto-main-base/src/main/java/com/facebook/presto/sql/planner/optimizations/ActualProperties.java
@@ -49,6 +49,7 @@
 import static com.google.common.base.MoreObjects.toStringHelper;
 import static com.google.common.base.Preconditions.checkArgument;
 import static com.google.common.collect.Iterables.transform;
+import static java.lang.String.format;
 import static java.util.Objects.requireNonNull;
 
 public class ActualProperties
@@ -56,11 +57,22 @@ public class ActualProperties
     private final Global global;
     private final List> localProperties;
     private final Map constants;
+    // Used to track the properties of the unique row_id
+    private final Optional propertiesFromUniqueColumn;
 
     private ActualProperties(
             Global global,
             List> localProperties,
             Map constants)
+    {
+        this(global, localProperties, constants, Optional.empty());
+    }
+
+    private ActualProperties(
+            Global global,
+            List> localProperties,
+            Map constants,
+            Optional propertiesFromUniqueColumn)
     {
         requireNonNull(global, "globalProperties is null");
         requireNonNull(localProperties, "localProperties is null");
@@ -85,6 +97,8 @@ private ActualProperties(
 
         this.localProperties = ImmutableList.copyOf(updatedLocalProperties);
         this.constants = ImmutableMap.copyOf(constants);
+        propertiesFromUniqueColumn.ifPresent(actualProperties -> checkArgument(!actualProperties.getPropertiesFromUniqueColumn().isPresent()));
+        this.propertiesFromUniqueColumn = propertiesFromUniqueColumn;
     }
 
     public boolean isCoordinatorOnly()
@@ -92,6 +106,11 @@ public boolean isCoordinatorOnly()
         return global.isCoordinatorOnly();
     }
 
+    public Optional getPropertiesFromUniqueColumn()
+    {
+        return propertiesFromUniqueColumn;
+    }
+
     /**
      * @return true if the plan will only execute on a single node
      */
@@ -120,6 +139,16 @@ public boolean isStreamPartitionedOn(Collection col
         }
     }
 
+    public boolean isStreamPartitionedOnAdditionalProperty(Collection columns, boolean exactly)
+    {
+        if (exactly) {
+            return propertiesFromUniqueColumn.isPresent() && propertiesFromUniqueColumn.get().global.isStreamPartitionedOnExactly(columns, ImmutableSet.of(), false);
+        }
+        else {
+            return propertiesFromUniqueColumn.isPresent() && propertiesFromUniqueColumn.get().global.isStreamPartitionedOn(columns, ImmutableSet.of(), false);
+        }
+    }
+
     public boolean isNodePartitionedOn(Collection columns, boolean exactly)
     {
         return isNodePartitionedOn(columns, false, exactly);
@@ -135,6 +164,16 @@ public boolean isNodePartitionedOn(Collection colum
         }
     }
 
+    public boolean isNodePartitionedOnAdditionalProperty(Collection columns, boolean exactly)
+    {
+        if (exactly) {
+            return propertiesFromUniqueColumn.isPresent() && propertiesFromUniqueColumn.get().global.isNodePartitionedOnExactly(columns, ImmutableSet.of(), false);
+        }
+        else {
+            return propertiesFromUniqueColumn.isPresent() && propertiesFromUniqueColumn.get().global.isNodePartitionedOn(columns, ImmutableSet.of(), false);
+        }
+    }
+
     @Deprecated
     public boolean isCompatibleTablePartitioningWith(Partitioning partitioning, boolean nullsAndAnyReplicated, Metadata metadata, Session session)
     {
@@ -194,6 +233,13 @@ public ActualProperties translateVariable(Function newAdditionalProperty = Optional.empty();
+        if (propertiesFromUniqueColumn.isPresent()) {
+            ActualProperties translatedAdditionalProperty = propertiesFromUniqueColumn.get().translateVariable(translator);
+            if (!translatedAdditionalProperty.getLocalProperties().isEmpty()) {
+                newAdditionalProperty = Optional.of(translatedAdditionalProperty);
+            }
+        }
         return builder()
                 .global(global.translateVariableToRowExpression(variable -> {
                     Optional translated = translator.apply(variable).map(RowExpression.class::cast);
@@ -204,6 +250,7 @@ public ActualProperties translateVariable(Function !inputToOutputVariables.containsKey(entry.getKey()))
                 .forEach(inputToOutputMappings::put);
+
+        Optional newAdditionalProperty = Optional.empty();
+        if (propertiesFromUniqueColumn.isPresent()) {
+            ActualProperties translatedAdditionalProperty = propertiesFromUniqueColumn.get().translateRowExpression(assignments);
+            if (!translatedAdditionalProperty.getLocalProperties().isEmpty()) {
+                newAdditionalProperty = Optional.of(translatedAdditionalProperty);
+            }
+        }
+
         return builder()
                 .global(global.translateRowExpression(inputToOutputMappings.build(), assignments))
                 .local(LocalProperties.translate(localProperties, variable -> Optional.ofNullable(inputToOutputVariables.get(variable))))
                 .constants(translatedConstants)
+                .propertiesFromUniqueColumn(newAdditionalProperty)
                 .build();
     }
 
@@ -274,6 +331,7 @@ public static class Builder
         private List> localProperties;
         private Map constants;
         private boolean unordered;
+        private Optional propertiesFromUniqueColumn;
 
         public Builder()
         {
@@ -281,10 +339,16 @@ public Builder()
         }
 
         public Builder(Global global, List> localProperties, Map constants)
+        {
+            this(global, localProperties, constants, Optional.empty());
+        }
+
+        public Builder(Global global, List> localProperties, Map constants, Optional propertiesFromUniqueColumn)
         {
             this.global = requireNonNull(global, "global is null");
             this.localProperties = ImmutableList.copyOf(localProperties);
             this.constants = ImmutableMap.copyOf(constants);
+            this.propertiesFromUniqueColumn = propertiesFromUniqueColumn;
         }
 
         public Builder global(Global global)
@@ -317,6 +381,18 @@ public Builder unordered(boolean unordered)
             return this;
         }
 
+        public Builder propertiesFromUniqueColumn(Optional propertiesFromUniqueColumn)
+        {
+            if (propertiesFromUniqueColumn.isPresent() && !propertiesFromUniqueColumn.get().getLocalProperties().isEmpty()) {
+                checkArgument(propertiesFromUniqueColumn.get().getLocalProperties().size() == 1);
+                this.propertiesFromUniqueColumn = propertiesFromUniqueColumn;
+            }
+            else {
+                this.propertiesFromUniqueColumn = Optional.empty();
+            }
+            return this;
+        }
+
         public ActualProperties build()
         {
             List> localProperties = this.localProperties;
@@ -332,14 +408,19 @@ public ActualProperties build()
                 }
                 localProperties = newLocalProperties.build();
             }
-            return new ActualProperties(global, localProperties, constants);
+            if (propertiesFromUniqueColumn.isPresent() && unordered) {
+                propertiesFromUniqueColumn = Optional.of(ActualProperties.builderFrom(propertiesFromUniqueColumn.get())
+                        .unordered(unordered)
+                        .build());
+            }
+            return new ActualProperties(global, localProperties, constants, propertiesFromUniqueColumn);
         }
     }
 
     @Override
     public int hashCode()
     {
-        return Objects.hash(global, localProperties, constants.keySet());
+        return Objects.hash(global, localProperties, constants.keySet(), propertiesFromUniqueColumn);
     }
 
     @Override
@@ -354,7 +435,8 @@ public boolean equals(Object obj)
         final ActualProperties other = (ActualProperties) obj;
         return Objects.equals(this.global, other.global)
                 && Objects.equals(this.localProperties, other.localProperties)
-                && Objects.equals(this.constants.keySet(), other.constants.keySet());
+                && Objects.equals(this.constants.keySet(), other.constants.keySet())
+                && Objects.equals(this.propertiesFromUniqueColumn, other.propertiesFromUniqueColumn);
     }
 
     @Override
@@ -364,6 +446,7 @@ public String toString()
                 .add("globalProperties", global)
                 .add("localProperties", localProperties)
                 .add("constants", constants)
+                .add("propertiesFromUniqueColumn", propertiesFromUniqueColumn)
                 .toString();
     }
 
@@ -391,7 +474,7 @@ private Global(Optional nodePartitioning, Optional s
                             || !streamPartitioning.isPresent()
                             || nodePartitioning.get().getVariableReferences().containsAll(streamPartitioning.get().getVariableReferences())
                             || streamPartitioning.get().getVariableReferences().containsAll(nodePartitioning.get().getVariableReferences()),
-                    "Global stream partitioning columns should match node partitioning columns");
+                    format("Global stream partitioning columns should match node partitioning columns, nodePartitioning: %s, streamPartitioning: %s", nodePartitioning, streamPartitioning));
             this.nodePartitioning = requireNonNull(nodePartitioning, "nodePartitioning is null");
             this.streamPartitioning = requireNonNull(streamPartitioning, "streamPartitioning is null");
             this.nullsAndAnyReplicated = nullsAndAnyReplicated;
diff --git a/presto-main-base/src/main/java/com/facebook/presto/sql/planner/optimizations/AddExchanges.java b/presto-main-base/src/main/java/com/facebook/presto/sql/planner/optimizations/AddExchanges.java
index 9749dd023c58e..9aba136f0f1c9 100644
--- a/presto-main-base/src/main/java/com/facebook/presto/sql/planner/optimizations/AddExchanges.java
+++ b/presto-main-base/src/main/java/com/facebook/presto/sql/planner/optimizations/AddExchanges.java
@@ -77,6 +77,7 @@
 import com.facebook.presto.sql.planner.plan.RowNumberNode;
 import com.facebook.presto.sql.planner.plan.SequenceNode;
 import com.facebook.presto.sql.planner.plan.StatisticsWriterNode;
+import com.facebook.presto.sql.planner.plan.TableFunctionNode;
 import com.facebook.presto.sql.planner.plan.TopNRowNumberNode;
 import com.google.common.annotations.VisibleForTesting;
 import com.google.common.cache.CacheBuilder;
@@ -311,7 +312,8 @@ public PlanWithProperties visitAggregation(AggregationNode node, PreferredProper
                         child.getProperties());
             }
             else if (hasMixedGroupingSets
-                    || !isStreamPartitionedOn(child.getProperties(), partitioningRequirement) && !isNodePartitionedOn(child.getProperties(), partitioningRequirement)) {
+                    || !isStreamPartitionedOn(child.getProperties(), partitioningRequirement) && !isNodePartitionedOn(child.getProperties(), partitioningRequirement)
+                    && !isNodePartitionedOnAdditionalProperty(child.getProperties(), partitioningRequirement) && !isStreamPartitionedOnAdditionalProperty(child.getProperties(), partitioningRequirement)) {
                 child = withDerivedProperties(
                         partitionedExchange(
                                 idAllocator.getNextId(),
@@ -410,6 +412,12 @@ public PlanWithProperties visitWindow(WindowNode node, PreferredProperties prefe
             return rebaseAndDeriveProperties(node, child);
         }
 
+        @Override
+        public PlanWithProperties visitTableFunction(TableFunctionNode node, PreferredProperties preferredProperties)
+        {
+            throw new UnsupportedOperationException("execution by operator is not yet implemented for table function " + node.getName());
+        }
+
         @Override
         public PlanWithProperties visitRowNumber(RowNumberNode node, PreferredProperties preferredProperties)
         {
@@ -1624,11 +1632,21 @@ private boolean isNodePartitionedOn(ActualProperties properties, Collection columns)
+        {
+            return properties.isNodePartitionedOnAdditionalProperty(columns, isExactPartitioningPreferred(session));
+        }
+
         private boolean isStreamPartitionedOn(ActualProperties properties, Collection columns)
         {
             return properties.isStreamPartitionedOn(columns, isExactPartitioningPreferred(session));
         }
 
+        private boolean isStreamPartitionedOnAdditionalProperty(ActualProperties properties, Collection columns)
+        {
+            return properties.isStreamPartitionedOnAdditionalProperty(columns, isExactPartitioningPreferred(session));
+        }
+
         private boolean shouldAggregationMergePartitionPreferences(AggregationPartitioningMergingStrategy aggregationPartitioningMergingStrategy)
         {
             if (isExactPartitioningPreferred(session)) {
diff --git a/presto-main-base/src/main/java/com/facebook/presto/sql/planner/optimizations/AddLocalExchanges.java b/presto-main-base/src/main/java/com/facebook/presto/sql/planner/optimizations/AddLocalExchanges.java
index 1618c69b543e6..65fa710226c48 100644
--- a/presto-main-base/src/main/java/com/facebook/presto/sql/planner/optimizations/AddLocalExchanges.java
+++ b/presto-main-base/src/main/java/com/facebook/presto/sql/planner/optimizations/AddLocalExchanges.java
@@ -399,7 +399,8 @@ public PlanWithProperties visitAggregation(AggregationNode node, StreamPreferred
             // [A, B] [(A, C)]     ->   List.of(Optional.of(GroupingProperty(C)))
             // [A, B] [(D, A, C)]  ->   List.of(Optional.of(GroupingProperty(D, C)))
             List>> matchResult = LocalProperties.match(child.getProperties().getLocalProperties(), LocalProperties.grouped(groupingKeys));
-            if (!matchResult.get(0).isPresent()) {
+            List>> matchResultForAdditional = LocalProperties.match(child.getProperties().getAdditionalLocalProperties(), LocalProperties.grouped(groupingKeys));
+            if (!matchResult.get(0).isPresent() || !matchResultForAdditional.get(0).isPresent()) {
                 // !isPresent() indicates the property was satisfied completely
                 preGroupedSymbols = groupingKeys;
             }
diff --git a/presto-main-base/src/main/java/com/facebook/presto/sql/planner/optimizations/KeyBasedSampler.java b/presto-main-base/src/main/java/com/facebook/presto/sql/planner/optimizations/KeyBasedSampler.java
index f34359aeba6af..1c7856a62caea 100644
--- a/presto-main-base/src/main/java/com/facebook/presto/sql/planner/optimizations/KeyBasedSampler.java
+++ b/presto-main-base/src/main/java/com/facebook/presto/sql/planner/optimizations/KeyBasedSampler.java
@@ -14,7 +14,6 @@
 package com.facebook.presto.sql.planner.optimizations;
 
 import com.facebook.presto.Session;
-import com.facebook.presto.common.QualifiedObjectName;
 import com.facebook.presto.common.function.OperatorType;
 import com.facebook.presto.common.type.Type;
 import com.facebook.presto.common.type.Varchars;
@@ -55,7 +54,6 @@
 import static com.facebook.presto.common.type.BooleanType.BOOLEAN;
 import static com.facebook.presto.common.type.DoubleType.DOUBLE;
 import static com.facebook.presto.common.type.VarcharType.VARCHAR;
-import static com.facebook.presto.metadata.BuiltInTypeAndFunctionNamespaceManager.JAVA_BUILTIN_NAMESPACE;
 import static com.facebook.presto.metadata.CastType.CAST;
 import static com.facebook.presto.spi.StandardErrorCode.FUNCTION_NOT_FOUND;
 import static com.facebook.presto.spi.StandardWarningCode.SAMPLED_FIELDS;
@@ -150,7 +148,7 @@ private PlanNode addSamplingFilter(PlanNode tableScanNode, Optional List> sorted(Collection columns, SortOrder
         return columns.stream().map(column -> new SortingProperty<>(column, order)).collect(toImmutableList());
     }
 
+    public static  List> unique(T column)
+    {
+        return ImmutableList.of(new UniqueProperty<>(column));
+    }
+
     public static  List> stripLeadingConstants(List> properties)
     {
         PeekingIterator> iterator = peekingIterator(properties.iterator());
diff --git a/presto-main-base/src/main/java/com/facebook/presto/sql/planner/optimizations/PropertyDerivations.java b/presto-main-base/src/main/java/com/facebook/presto/sql/planner/optimizations/PropertyDerivations.java
index f086283a140da..bb4f27b09bb0e 100644
--- a/presto-main-base/src/main/java/com/facebook/presto/sql/planner/optimizations/PropertyDerivations.java
+++ b/presto-main-base/src/main/java/com/facebook/presto/sql/planner/optimizations/PropertyDerivations.java
@@ -23,6 +23,7 @@
 import com.facebook.presto.spi.GroupingProperty;
 import com.facebook.presto.spi.LocalProperty;
 import com.facebook.presto.spi.SortingProperty;
+import com.facebook.presto.spi.UniqueProperty;
 import com.facebook.presto.spi.plan.AggregationNode;
 import com.facebook.presto.spi.plan.DeleteNode;
 import com.facebook.presto.spi.plan.DistinctLimitNode;
@@ -90,6 +91,7 @@
 
 import static com.facebook.presto.SystemSessionProperties.isJoinSpillingEnabled;
 import static com.facebook.presto.SystemSessionProperties.isSpillEnabled;
+import static com.facebook.presto.SystemSessionProperties.isUtilizeUniquePropertyInQueryPlanningEnabled;
 import static com.facebook.presto.SystemSessionProperties.planWithTableNodePartitioning;
 import static com.facebook.presto.common.predicate.TupleDomain.toLinkedMap;
 import static com.facebook.presto.spi.relation.DomainTranslator.BASIC_COLUMN_EXTRACTOR;
@@ -142,6 +144,22 @@ public static ActualProperties streamBackdoorDeriveProperties(PlanNode node, Lis
         return node.accept(new Visitor(metadata, session), inputProperties);
     }
 
+    public static Optional uniqueToGroupProperties(ActualProperties properties)
+    {
+        // We only call uniqueToGroupProperties on derived properties from propertiesFromUniqueColumn, which can have one local property if the column is preserved in a node
+        // output, or no local property if the column is not preserved in a node output
+        checkArgument(properties.getLocalProperties().size() <= 1);
+        if (properties.getLocalProperties().isEmpty()) {
+            return Optional.empty();
+        }
+        LocalProperty localProperty = Iterables.getOnlyElement(properties.getLocalProperties());
+        if (localProperty instanceof UniqueProperty) {
+            return Optional.of(ActualProperties.builderFrom(properties).local(ImmutableList.of(new GroupingProperty<>(ImmutableList.of(((UniqueProperty) localProperty).getColumn())))).build());
+        }
+        checkState(localProperty instanceof GroupingProperty, "returned actual properties should have grouping property");
+        return Optional.of(properties);
+    }
+
     private static class Visitor
             extends InternalPlanVisitor>
     {
@@ -196,12 +214,14 @@ public ActualProperties visitAssignUniqueId(AssignUniqueId node, List in
                 inputToOutputMappings.putIfAbsent(argument, argument);
             }
 
-            return Iterables.getOnlyElement(inputProperties).translateVariable(column -> Optional.ofNullable(inputToOutputMappings.get(column)));
+            ActualProperties properties = Iterables.getOnlyElement(inputProperties);
+            return ActualProperties.builderFrom(properties.translateVariable(column -> Optional.ofNullable(inputToOutputMappings.get(column))))
+                    .propertiesFromUniqueColumn(properties.getPropertiesFromUniqueColumn().flatMap(x -> uniqueToGroupProperties(x.translateVariable(column -> Optional.ofNullable(inputToOutputMappings.get(column))))))
+                    .build();
         }
 
         @Override
@@ -291,10 +314,9 @@ public ActualProperties visitAggregation(AggregationNode node, List node.getGroupingKeys().contains(variable) ? Optional.of(variable) : Optional.empty());
-
             return ActualProperties.builderFrom(translated)
                     .local(LocalProperties.grouped(node.getGroupingKeys()))
-                    .build();
+                    .propertiesFromUniqueColumn(uniqueProperties(translated.getPropertiesFromUniqueColumn())).build();
         }
 
         @Override
@@ -303,6 +325,14 @@ public ActualProperties visitRowNumber(RowNumberNode node, List uniqueProperties(Optional properties)
+        {
+            if (properties.isPresent() && properties.get().getLocalProperties().size() == 1 && properties.get().getLocalProperties().get(0) instanceof UniqueProperty) {
+                return properties;
+            }
+            return Optional.empty();
+        }
+
         @Override
         public ActualProperties visitTopNRowNumber(TopNRowNumberNode node, List inputProperties)
         {
@@ -316,6 +346,7 @@ public ActualProperties visitTopNRowNumber(TopNRowNumberNode node, List inputPro
 
             return ActualProperties.builderFrom(properties)
                     .local(localProperties)
+                    .propertiesFromUniqueColumn(uniqueProperties(properties.getPropertiesFromUniqueColumn()))
                     .build();
         }
 
@@ -344,6 +376,7 @@ public ActualProperties visitSort(SortNode node, List inputPro
 
             return ActualProperties.builderFrom(properties)
                     .local(localProperties)
+                    .propertiesFromUniqueColumn(uniqueProperties(properties.getPropertiesFromUniqueColumn()))
                     .build();
         }
 
@@ -360,6 +393,7 @@ public ActualProperties visitDistinctLimit(DistinctLimitNode node, List inputPro
 
                     return ActualProperties.builderFrom(probeProperties)
                             .constants(constants)
+                            .propertiesFromUniqueColumn(probeProperties.getPropertiesFromUniqueColumn().flatMap(x -> uniqueToGroupProperties(x.translateVariable(column -> filterOrRewrite(outputVariableReferences, node.getCriteria(), column)))))
                             .unordered(unordered)
                             .build();
                 case LEFT:
                     return ActualProperties.builderFrom(probeProperties.translateVariable(column -> filterIfMissing(outputVariableReferences, column)))
+                            .propertiesFromUniqueColumn(probeProperties.getPropertiesFromUniqueColumn().flatMap(x -> uniqueToGroupProperties(x.translateVariable(column -> filterIfMissing(outputVariableReferences, column)))))
                             .unordered(unordered)
                             .build();
                 case RIGHT:
@@ -518,10 +554,12 @@ public ActualProperties visitIndexJoin(IndexJoinNode node, List uniqueToGroupProperties(x)))
                             .build();
                 case SOURCE_OUTER:
                     return ActualProperties.builderFrom(probeProperties)
                             .constants(probeProperties.getConstants())
+                            .propertiesFromUniqueColumn(probeProperties.getPropertiesFromUniqueColumn().flatMap(x -> uniqueToGroupProperties(x)))
                             .build();
                 default:
                     throw new UnsupportedOperationException("Unsupported join type: " + node.getType());
@@ -553,16 +591,19 @@ public ActualProperties visitMergeJoin(MergeJoinNode node, List uniqueToGroupProperties(x.translateVariable(column -> filterOrRewrite(outputVariableReferences, node.getCriteria(), column)))))
                             .constants(constants)
                             .build();
                 case LEFT:
                     return ActualProperties.builderFrom(leftProperties.translateVariable(column -> filterIfMissing(outputVariableReferences, column)))
+                            .propertiesFromUniqueColumn(leftProperties.getPropertiesFromUniqueColumn().flatMap(x -> uniqueToGroupProperties(x.translateVariable(column -> filterIfMissing(outputVariableReferences, column)))))
                             .build();
                 case RIGHT:
                     rightProperties = rightProperties.translateVariable(column -> filterIfMissing(node.getOutputVariables(), column));
 
                     return ActualProperties.builderFrom(rightProperties.translateVariable(column -> filterIfMissing(outputVariableReferences, column)))
                             .local(ImmutableList.of())
+                            .propertiesFromUniqueColumn(rightProperties.getPropertiesFromUniqueColumn().flatMap(x -> uniqueToGroupProperties(x.translateVariable(column -> filterIfMissing(outputVariableReferences, column)))))
                             .unordered(true)
                             .build();
                 case FULL:
@@ -598,6 +639,7 @@ public ActualProperties visitExchange(ExchangeNode node, List
             checkArgument(!node.getScope().isRemote() || inputProperties.stream().noneMatch(ActualProperties::isNullsAndAnyReplicated), "Null-and-any replicated inputs should not be remotely exchanged");
 
             Set> entries = null;
+            ActualProperties translated = null;
             for (int sourceIndex = 0; sourceIndex < node.getSources().size(); sourceIndex++) {
                 List inputVariables = node.getInputs().get(sourceIndex);
                 Map inputToOutput = new HashMap<>();
@@ -605,7 +647,7 @@ public ActualProperties visitExchange(ExchangeNode node, List
                     inputToOutput.put(inputVariables.get(i), node.getOutputVariables().get(i));
                 }
 
-                ActualProperties translated = inputProperties.get(sourceIndex).translateVariable(variable -> Optional.ofNullable(inputToOutput.get(variable)));
+                translated = inputProperties.get(sourceIndex).translateVariable(variable -> Optional.ofNullable(inputToOutput.get(variable)));
 
                 entries = (entries == null) ? translated.getConstants().entrySet() : Sets.intersection(entries, translated.getConstants().entrySet());
             }
@@ -621,6 +663,8 @@ public ActualProperties visitExchange(ExchangeNode node, List
                         .forEach(localProperties::add);
             }
 
+            boolean additionalPropertyIsUnique = inputProperties.size() == 1 && uniqueProperties(translated.getPropertiesFromUniqueColumn()).isPresent() && !node.getType().equals(ExchangeNode.Type.REPLICATE);
+
             // Local exchanges are only created in AddLocalExchanges, at the end of optimization, and
             // local exchanges do not produce all global properties as represented by ActualProperties.
             // This is acceptable because AddLocalExchanges does not use global properties and is only
@@ -640,6 +684,10 @@ else if (inputProperties.stream().anyMatch(ActualProperties::isSingleNode)) {
                     builder.global(coordinatorSingleStreamPartition());
                 }
 
+                if (additionalPropertyIsUnique) {
+                    builder.propertiesFromUniqueColumn(translated.getPropertiesFromUniqueColumn());
+                }
+
                 return builder.build();
             }
 
@@ -650,6 +698,7 @@ else if (inputProperties.stream().anyMatch(ActualProperties::isSingleNode)) {
                             .global(coordinatorOnly ? coordinatorSingleStreamPartition() : singleStreamPartition())
                             .local(localProperties.build())
                             .constants(constants)
+                            .propertiesFromUniqueColumn(additionalPropertyIsUnique ? translated.getPropertiesFromUniqueColumn() : Optional.empty())
                             .build();
                 case REPARTITION: {
                     Global globalPartitioning;
@@ -666,6 +715,7 @@ else if (inputProperties.stream().anyMatch(ActualProperties::isSingleNode)) {
                     return ActualProperties.builder()
                             .global(globalPartitioning)
                             .constants(constants)
+                            .propertiesFromUniqueColumn(additionalPropertyIsUnique ? translated.getPropertiesFromUniqueColumn() : Optional.empty())
                             .build();
                 }
                 case REPLICATE:
@@ -691,6 +741,7 @@ public ActualProperties visitFilter(FilterNode node, List inpu
 
             return ActualProperties.builderFrom(properties)
                     .constants(constants)
+                    .propertiesFromUniqueColumn(properties.getPropertiesFromUniqueColumn())
                     .build();
         }
 
@@ -728,6 +779,7 @@ else if (!(value instanceof RowExpression)) {
 
             return ActualProperties.builderFrom(translatedProperties)
                     .constants(constants)
+                    .propertiesFromUniqueColumn(properties.getPropertiesFromUniqueColumn().map(x -> x.translateRowExpression(node.getAssignments().getMap())))
                     .build();
         }
 
@@ -814,6 +866,13 @@ public ActualProperties visitTableScan(TableScanNode node, List Optional.ofNullable(assignments.get(column))));
+            if (isUtilizeUniquePropertyInQueryPlanningEnabled(session) && layout.getUniqueColumn().isPresent() && assignments.containsKey(layout.getUniqueColumn().get())) {
+                VariableReferenceExpression uniqueVariable = assignments.get(layout.getUniqueColumn().get());
+                ActualProperties.Builder propertiesFromUniqueColumn = ActualProperties.builder();
+                propertiesFromUniqueColumn.global(partitionedOn(ARBITRARY_DISTRIBUTION, ImmutableList.of(uniqueVariable), Optional.empty()));
+                propertiesFromUniqueColumn.local(LocalProperties.unique(uniqueVariable));
+                properties.propertiesFromUniqueColumn(Optional.of(propertiesFromUniqueColumn.build()));
+            }
 
             return properties.build();
         }
diff --git a/presto-main-base/src/main/java/com/facebook/presto/sql/planner/optimizations/StreamPreferredProperties.java b/presto-main-base/src/main/java/com/facebook/presto/sql/planner/optimizations/StreamPreferredProperties.java
index 221d224ea6d40..96288e6af729c 100644
--- a/presto-main-base/src/main/java/com/facebook/presto/sql/planner/optimizations/StreamPreferredProperties.java
+++ b/presto-main-base/src/main/java/com/facebook/presto/sql/planner/optimizations/StreamPreferredProperties.java
@@ -188,9 +188,11 @@ else if (actualProperties.getDistribution() == SINGLE) {
         // is there a preference for a specific partitioning scheme?
         if (partitioningColumns.isPresent()) {
             if (exactColumnOrder) {
-                return actualProperties.isExactlyPartitionedOn(partitioningColumns.get());
+                return actualProperties.isExactlyPartitionedOn(partitioningColumns.get())
+                        || actualProperties.getStreamPropertiesFromUniqueColumn().isPresent() && actualProperties.getStreamPropertiesFromUniqueColumn().get().isExactlyPartitionedOn(partitioningColumns.get());
             }
-            return actualProperties.isPartitionedOn(partitioningColumns.get());
+            return actualProperties.isPartitionedOn(partitioningColumns.get())
+                    || actualProperties.getStreamPropertiesFromUniqueColumn().isPresent() && actualProperties.getStreamPropertiesFromUniqueColumn().get().isPartitionedOn(partitioningColumns.get());
         }
 
         return true;
diff --git a/presto-main-base/src/main/java/com/facebook/presto/sql/planner/optimizations/StreamPropertyDerivations.java b/presto-main-base/src/main/java/com/facebook/presto/sql/planner/optimizations/StreamPropertyDerivations.java
index 58f8e32977de6..51e6d3de35ab2 100644
--- a/presto-main-base/src/main/java/com/facebook/presto/sql/planner/optimizations/StreamPropertyDerivations.java
+++ b/presto-main-base/src/main/java/com/facebook/presto/sql/planner/optimizations/StreamPropertyDerivations.java
@@ -18,6 +18,7 @@
 import com.facebook.presto.metadata.TableLayout;
 import com.facebook.presto.spi.ColumnHandle;
 import com.facebook.presto.spi.LocalProperty;
+import com.facebook.presto.spi.UniqueProperty;
 import com.facebook.presto.spi.plan.AggregationNode;
 import com.facebook.presto.spi.plan.DeleteNode;
 import com.facebook.presto.spi.plan.DistinctLimitNode;
@@ -78,6 +79,7 @@
 import java.util.function.Function;
 import java.util.stream.Collectors;
 
+import static com.facebook.presto.SystemSessionProperties.isUtilizeUniquePropertyInQueryPlanningEnabled;
 import static com.facebook.presto.sql.planner.SystemPartitioningHandle.FIXED_ARBITRARY_DISTRIBUTION;
 import static com.facebook.presto.sql.planner.optimizations.PropertyDerivations.extractFixedValuesToConstantExpressions;
 import static com.facebook.presto.sql.planner.optimizations.StreamPropertyDerivations.StreamProperties.StreamDistribution.FIXED;
@@ -123,7 +125,10 @@ public static StreamProperties deriveProperties(PlanNode node, List properties.otherActualProperties)
+                        .map(properties -> {
+                            checkState(properties.otherActualProperties.isPresent(), "otherActualProperties should always exist");
+                            return properties.otherActualProperties.get();
+                        })
                         .collect(toImmutableList()),
                 metadata,
                 session);
@@ -225,11 +230,8 @@ public StreamProperties visitIndexJoin(IndexJoinNode node, List streamPropertiesFromUniqueColumn = Optional.empty();
+            if (isUtilizeUniquePropertyInQueryPlanningEnabled(session) && layout.getUniqueColumn().isPresent() && assignments.containsKey(layout.getUniqueColumn().get())) {
+                streamPropertiesFromUniqueColumn = Optional.of(new StreamProperties(streamDistribution, Optional.of(ImmutableList.of(assignments.get(layout.getUniqueColumn().get()))), false));
+            }
+
             // if we are partitioned on empty set, we must say multiple of unknown partitioning, because
             // the connector does not guarantee a single split in this case (since it might not understand
             // that the value is a constant).
             if (streamPartitionSymbols.isPresent() && streamPartitionSymbols.get().isEmpty()) {
-                return new StreamProperties(streamDistribution, Optional.empty(), false);
+                return new StreamProperties(streamDistribution, Optional.empty(), false, Optional.empty(), streamPropertiesFromUniqueColumn);
             }
-            return new StreamProperties(streamDistribution, streamPartitionSymbols, false);
+            return new StreamProperties(streamDistribution, streamPartitionSymbols, false, Optional.empty(), streamPropertiesFromUniqueColumn);
         }
 
         private Optional> getNonConstantVariables(Set columnHandles, Map assignments, Set globalConstants)
@@ -337,20 +344,32 @@ public StreamProperties visitExchange(ExchangeNode node, List
                 return new StreamProperties(MULTIPLE, Optional.empty(), false);
             }
 
+            Optional additionalUniqueProperty = Optional.empty();
+            if (inputProperties.size() == 1 && inputProperties.get(0).hasUniqueProperties() && !node.getType().equals(ExchangeNode.Type.REPLICATE)) {
+                List inputVariables = node.getInputs().get(0);
+                Map inputToOutput = new HashMap<>();
+                for (int i = 0; i < node.getOutputVariables().size(); i++) {
+                    inputToOutput.put(inputVariables.get(i), node.getOutputVariables().get(i));
+                }
+                checkState(inputProperties.get(0).getStreamPropertiesFromUniqueColumn().isPresent(),
+                        "when unique columns exists, the stream is also partitioned by the unique column and should be represented in the streamPropertiesFromUniqueColumn field");
+                additionalUniqueProperty = Optional.of(inputProperties.get(0).getStreamPropertiesFromUniqueColumn().get().translate(column -> Optional.ofNullable(inputToOutput.get(column))));
+            }
+
             if (node.getScope().isRemote()) {
                 // TODO: correctly determine if stream is parallelised
                 // based on session properties
-                return StreamProperties.fixedStreams();
+                return StreamProperties.fixedStreams().withStreamPropertiesFromUniqueColumn(additionalUniqueProperty);
             }
 
             switch (node.getType()) {
                 case GATHER:
-                    return StreamProperties.singleStream();
+                    return StreamProperties.singleStream().withStreamPropertiesFromUniqueColumn(additionalUniqueProperty);
                 case REPARTITION:
                     if (node.getPartitioningScheme().getPartitioning().getHandle().equals(FIXED_ARBITRARY_DISTRIBUTION) ||
                             // no strict partitioning guarantees when multiple writers per partitions are allows (scaled writers)
                             node.getPartitioningScheme().isScaleWriters()) {
-                        return new StreamProperties(FIXED, Optional.empty(), false);
+                        return new StreamProperties(FIXED, Optional.empty(), false).withStreamPropertiesFromUniqueColumn(additionalUniqueProperty);
                     }
                     checkArgument(
                             node.getPartitioningScheme().getPartitioning().getArguments().stream().allMatch(VariableReferenceExpression.class::isInstance),
@@ -359,7 +378,7 @@ public StreamProperties visitExchange(ExchangeNode node, List
                             FIXED,
                             Optional.of(node.getPartitioningScheme().getPartitioning().getArguments().stream()
                                     .map(VariableReferenceExpression.class::cast)
-                                    .collect(toImmutableList())), false);
+                                    .collect(toImmutableList())), false).withStreamPropertiesFromUniqueColumn(additionalUniqueProperty);
                 case REPLICATE:
                     return new StreamProperties(MULTIPLE, Optional.empty(), false);
             }
@@ -668,21 +687,33 @@ public enum StreamDistribution
 
         // We are only interested in the local properties, but PropertyDerivations requires input
         // ActualProperties, so we hold on to the whole object
-        private final ActualProperties otherActualProperties;
+        private final Optional otherActualProperties;
 
         // NOTE: Partitioning on zero columns (or effectively zero columns if the columns are constant) indicates that all
         // the rows will be partitioned into a single stream.
 
+        private final Optional streamPropertiesFromUniqueColumn;
+
         private StreamProperties(StreamDistribution distribution, Optional> partitioningColumns, boolean ordered)
         {
-            this(distribution, partitioningColumns, ordered, null);
+            this(distribution, partitioningColumns, ordered, Optional.empty());
+        }
+
+        private StreamProperties(
+                StreamDistribution distribution,
+                Optional> partitioningColumns,
+                boolean ordered,
+                Optional otherActualProperties)
+        {
+            this(distribution, partitioningColumns, ordered, otherActualProperties, Optional.empty());
         }
 
         private StreamProperties(
                 StreamDistribution distribution,
                 Optional> partitioningColumns,
                 boolean ordered,
-                ActualProperties otherActualProperties)
+                Optional otherActualProperties,
+                Optional streamPropertiesFromUniqueColumn)
         {
             this.distribution = requireNonNull(distribution, "distribution is null");
 
@@ -697,13 +728,38 @@ private StreamProperties(
             this.ordered = ordered;
             checkArgument(!ordered || distribution == SINGLE, "Ordered must be a single stream");
 
-            this.otherActualProperties = otherActualProperties;
+            this.otherActualProperties = requireNonNull(otherActualProperties);
+            requireNonNull(streamPropertiesFromUniqueColumn).ifPresent(properties -> checkArgument(!properties.streamPropertiesFromUniqueColumn.isPresent()));
+            // When unique properties exists, the stream is also partitioned on the unique column
+            if (otherActualProperties.isPresent() && otherActualProperties.get().getPropertiesFromUniqueColumn().isPresent()) {
+                ActualProperties propertiesFromUniqueColumn = otherActualProperties.get().getPropertiesFromUniqueColumn().get();
+                if (Iterables.getOnlyElement(propertiesFromUniqueColumn.getLocalProperties()) instanceof UniqueProperty) {
+                    VariableReferenceExpression uniqueVariable = (VariableReferenceExpression) ((UniqueProperty) Iterables.getOnlyElement(propertiesFromUniqueColumn.getLocalProperties())).getColumn();
+                    checkState(streamPropertiesFromUniqueColumn.isPresent() && streamPropertiesFromUniqueColumn.get().partitioningColumns.isPresent()
+                            && Iterables.getOnlyElement(streamPropertiesFromUniqueColumn.get().partitioningColumns.get()).equals(uniqueVariable));
+                }
+            }
+            this.streamPropertiesFromUniqueColumn = streamPropertiesFromUniqueColumn;
         }
 
         public List> getLocalProperties()
         {
-            checkState(otherActualProperties != null, "otherActualProperties not set");
-            return otherActualProperties.getLocalProperties();
+            checkState(otherActualProperties.isPresent(), "otherActualProperties not set");
+            return otherActualProperties.get().getLocalProperties();
+        }
+
+        public List> getAdditionalLocalProperties()
+        {
+            checkState(otherActualProperties.isPresent(), "otherActualProperties not set");
+            if (!otherActualProperties.get().getPropertiesFromUniqueColumn().isPresent()) {
+                return ImmutableList.of();
+            }
+            return otherActualProperties.get().getPropertiesFromUniqueColumn().get().getLocalProperties();
+        }
+
+        public Optional getStreamPropertiesFromUniqueColumn()
+        {
+            return streamPropertiesFromUniqueColumn;
         }
 
         private static StreamProperties singleStream()
@@ -725,8 +781,9 @@ private StreamProperties unordered(boolean unordered)
         {
             if (unordered) {
                 ActualProperties updatedProperties = null;
-                if (otherActualProperties != null) {
-                    updatedProperties = ActualProperties.builderFrom(otherActualProperties)
+                if (otherActualProperties.isPresent()) {
+                    updatedProperties = ActualProperties.builderFrom(otherActualProperties.get())
+                            .propertiesFromUniqueColumn(otherActualProperties.get().getPropertiesFromUniqueColumn().map(x -> ActualProperties.builderFrom(x).unordered(true).build()))
                             .unordered(true)
                             .build();
                 }
@@ -734,11 +791,33 @@ private StreamProperties unordered(boolean unordered)
                         distribution,
                         partitioningColumns,
                         false,
-                        updatedProperties);
+                        Optional.ofNullable(updatedProperties),
+                        streamPropertiesFromUniqueColumn.map(x -> x.unordered(true)));
             }
             return this;
         }
 
+        public StreamProperties uniqueToGroupProperties()
+        {
+            if (otherActualProperties.isPresent() && otherActualProperties.get().getPropertiesFromUniqueColumn().isPresent()) {
+                if (Iterables.getOnlyElement(otherActualProperties.get().getPropertiesFromUniqueColumn().get().getLocalProperties()) instanceof UniqueProperty) {
+                    Optional groupedProperties = PropertyDerivations.uniqueToGroupProperties(otherActualProperties.get().getPropertiesFromUniqueColumn().get());
+                    return new StreamProperties(distribution, partitioningColumns, ordered,
+                            otherActualProperties.map(x -> ActualProperties.builderFrom(x).propertiesFromUniqueColumn(groupedProperties).build()),
+                            streamPropertiesFromUniqueColumn.map(StreamProperties::uniqueToGroupProperties));
+                }
+            }
+            return this;
+        }
+
+        public boolean hasUniqueProperties()
+        {
+            if (otherActualProperties.isPresent() && otherActualProperties.get().getPropertiesFromUniqueColumn().isPresent()) {
+                return Iterables.getOnlyElement(otherActualProperties.get().getPropertiesFromUniqueColumn().get().getLocalProperties()) instanceof UniqueProperty;
+            }
+            return false;
+        }
+
         public boolean isSingleStream()
         {
             return distribution == SINGLE;
@@ -783,7 +862,12 @@ private StreamProperties withUnspecifiedPartitioning()
 
         private StreamProperties withOtherActualProperties(ActualProperties actualProperties)
         {
-            return new StreamProperties(distribution, partitioningColumns, ordered, actualProperties);
+            return new StreamProperties(distribution, partitioningColumns, ordered, Optional.ofNullable(actualProperties), streamPropertiesFromUniqueColumn);
+        }
+
+        private StreamProperties withStreamPropertiesFromUniqueColumn(Optional streamPropertiesFromUniqueColumn)
+        {
+            return new StreamProperties(distribution, partitioningColumns, ordered, otherActualProperties, streamPropertiesFromUniqueColumn);
         }
 
         public StreamProperties translate(Function> translator)
@@ -801,7 +885,8 @@ public StreamProperties translate(Function x.translateVariable(translator)),
+                    streamPropertiesFromUniqueColumn.map(x -> x.translate(translator)));
         }
 
         public Optional> getPartitioningColumns()
@@ -812,7 +897,7 @@ public Optional> getPartitioningColumns()
         @Override
         public int hashCode()
         {
-            return Objects.hash(distribution, partitioningColumns);
+            return Objects.hash(distribution, partitioningColumns, ordered, otherActualProperties, streamPropertiesFromUniqueColumn);
         }
 
         @Override
@@ -826,7 +911,10 @@ public boolean equals(Object obj)
             }
             StreamProperties other = (StreamProperties) obj;
             return Objects.equals(this.distribution, other.distribution) &&
-                    Objects.equals(this.partitioningColumns, other.partitioningColumns);
+                    Objects.equals(this.partitioningColumns, other.partitioningColumns) &&
+                    this.ordered == other.ordered &&
+                    Objects.equals(this.otherActualProperties, other.otherActualProperties) &&
+                    Objects.equals(this.streamPropertiesFromUniqueColumn, other.streamPropertiesFromUniqueColumn);
         }
 
         @Override
@@ -835,6 +923,9 @@ public String toString()
             return toStringHelper(this)
                     .add("distribution", distribution)
                     .add("partitioningColumns", partitioningColumns)
+                    .add("ordered", ordered)
+                    .add("otherActualProperties", otherActualProperties)
+                    .add("streamPropertiesFromUniqueColumn", streamPropertiesFromUniqueColumn)
                     .toString();
         }
     }
diff --git a/presto-main-base/src/main/java/com/facebook/presto/sql/planner/optimizations/UnaliasSymbolReferences.java b/presto-main-base/src/main/java/com/facebook/presto/sql/planner/optimizations/UnaliasSymbolReferences.java
index eb61dd1bf42cb..c279a0675ec7f 100644
--- a/presto-main-base/src/main/java/com/facebook/presto/sql/planner/optimizations/UnaliasSymbolReferences.java
+++ b/presto-main-base/src/main/java/com/facebook/presto/sql/planner/optimizations/UnaliasSymbolReferences.java
@@ -74,6 +74,7 @@
 import com.facebook.presto.sql.planner.plan.SequenceNode;
 import com.facebook.presto.sql.planner.plan.SimplePlanRewriter;
 import com.facebook.presto.sql.planner.plan.StatisticsWriterNode;
+import com.facebook.presto.sql.planner.plan.TableFunctionNode;
 import com.facebook.presto.sql.planner.plan.TableWriterMergeNode;
 import com.facebook.presto.sql.planner.plan.TopNRowNumberNode;
 import com.facebook.presto.sql.planner.plan.UpdateNode;
@@ -477,6 +478,21 @@ public PlanNode visitTableFinish(TableFinishNode node, RewriteContext cont
             return mapper.map(node, source);
         }
 
+        @Override
+        public PlanNode visitTableFunction(TableFunctionNode node, RewriteContext context)
+        {
+            return new TableFunctionNode(
+                    node.getSourceLocation(),
+                    node.getId(),
+                    Optional.empty(),
+                    node.getName(),
+                    node.getArguments(),
+                    node.getOutputVariables(),
+                    node.getSources(),
+                    node.getTableArgumentProperties(),
+                    node.getHandle());
+        }
+
         @Override
         public PlanNode visitRowNumber(RowNumberNode node, RewriteContext context)
         {
diff --git a/presto-main-base/src/main/java/com/facebook/presto/sql/planner/plan/InternalPlanVisitor.java b/presto-main-base/src/main/java/com/facebook/presto/sql/planner/plan/InternalPlanVisitor.java
index 15b53e87ac787..58ef334a0f01e 100644
--- a/presto-main-base/src/main/java/com/facebook/presto/sql/planner/plan/InternalPlanVisitor.java
+++ b/presto-main-base/src/main/java/com/facebook/presto/sql/planner/plan/InternalPlanVisitor.java
@@ -126,4 +126,9 @@ public R visitSequence(SequenceNode node, C context)
     {
         return visitPlan(node, context);
     }
+
+    public R visitTableFunction(TableFunctionNode node, C context)
+    {
+        return visitPlan(node, context);
+    }
 }
diff --git a/presto-main-base/src/main/java/com/facebook/presto/sql/planner/plan/Patterns.java b/presto-main-base/src/main/java/com/facebook/presto/sql/planner/plan/Patterns.java
index a2d18830e9d90..f1a00c4b1a128 100644
--- a/presto-main-base/src/main/java/com/facebook/presto/sql/planner/plan/Patterns.java
+++ b/presto-main-base/src/main/java/com/facebook/presto/sql/planner/plan/Patterns.java
@@ -229,6 +229,11 @@ public static Pattern window()
         return typeOf(WindowNode.class);
     }
 
+    public static Pattern tableFunction()
+    {
+        return typeOf(TableFunctionNode.class);
+    }
+
     public static Pattern rowNumber()
     {
         return typeOf(RowNumberNode.class);
diff --git a/presto-main-base/src/main/java/com/facebook/presto/sql/planner/plan/TableFunctionNode.java b/presto-main-base/src/main/java/com/facebook/presto/sql/planner/plan/TableFunctionNode.java
new file mode 100644
index 0000000000000..97892523498c0
--- /dev/null
+++ b/presto-main-base/src/main/java/com/facebook/presto/sql/planner/plan/TableFunctionNode.java
@@ -0,0 +1,178 @@
+/*
+ * Licensed 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 com.facebook.presto.sql.planner.plan;
+
+import com.facebook.presto.metadata.TableFunctionHandle;
+import com.facebook.presto.spi.SourceLocation;
+import com.facebook.presto.spi.function.table.Argument;
+import com.facebook.presto.spi.plan.DataOrganizationSpecification;
+import com.facebook.presto.spi.plan.PlanNode;
+import com.facebook.presto.spi.plan.PlanNodeId;
+import com.facebook.presto.spi.relation.VariableReferenceExpression;
+import com.fasterxml.jackson.annotation.JsonCreator;
+import com.fasterxml.jackson.annotation.JsonProperty;
+import com.google.errorprone.annotations.Immutable;
+
+import java.util.List;
+import java.util.Map;
+import java.util.Optional;
+
+import static com.google.common.base.Preconditions.checkArgument;
+import static java.util.Objects.requireNonNull;
+
+@Immutable
+public class TableFunctionNode
+        extends InternalPlanNode
+{
+    private final String name;
+    private final Map arguments;
+    private final List outputVariables;
+    private final List sources;
+    private final List tableArgumentProperties;
+    private final TableFunctionHandle handle;
+
+    @JsonCreator
+    public TableFunctionNode(
+            @JsonProperty("id") PlanNodeId id,
+            @JsonProperty("name") String name,
+            @JsonProperty("arguments") Map arguments,
+            @JsonProperty("outputVariables") List outputVariables,
+            @JsonProperty("sources") List sources,
+            @JsonProperty("tableArgumentProperties") List tableArgumentProperties,
+            @JsonProperty("handle") TableFunctionHandle handle)
+    {
+        this(Optional.empty(), id, Optional.empty(), name, arguments, outputVariables, sources, tableArgumentProperties, handle);
+    }
+
+    public TableFunctionNode(
+            Optional sourceLocation,
+            PlanNodeId id,
+            Optional statsEquivalentPlanNode,
+            String name,
+            Map arguments,
+            List outputVariables,
+            List sources,
+            List tableArgumentProperties,
+            TableFunctionHandle handle)
+    {
+        super(sourceLocation, id, statsEquivalentPlanNode);
+        this.name = requireNonNull(name, "name is null");
+        this.arguments = requireNonNull(arguments, "arguments is null");
+        this.outputVariables = requireNonNull(outputVariables, "outputVariables is null");
+        this.sources = requireNonNull(sources, "sources is null");
+        this.tableArgumentProperties = requireNonNull(tableArgumentProperties, "tableArgumentProperties is null");
+        this.handle = requireNonNull(handle, "handle is null");
+    }
+
+    @JsonProperty
+    public String getName()
+    {
+        return name;
+    }
+
+    @JsonProperty
+    public Map getArguments()
+    {
+        return arguments;
+    }
+
+    @JsonProperty
+    public List getOutputVariables()
+    {
+        return outputVariables;
+    }
+
+    @JsonProperty
+    public List getTableArgumentProperties()
+    {
+        return tableArgumentProperties;
+    }
+
+    @JsonProperty
+    public TableFunctionHandle getHandle()
+    {
+        return handle;
+    }
+
+    @JsonProperty
+    @Override
+    public List getSources()
+    {
+        return sources;
+    }
+
+    @Override
+    public  R accept(InternalPlanVisitor visitor, C context)
+    {
+        return visitor.visitTableFunction(this, context);
+    }
+
+    @Override
+    public PlanNode replaceChildren(List newSources)
+    {
+        checkArgument(sources.size() == newSources.size(), "wrong number of new children");
+        return new TableFunctionNode(getId(), name, arguments, outputVariables, newSources, tableArgumentProperties, handle);
+    }
+
+    @Override
+    public PlanNode assignStatsEquivalentPlanNode(Optional statsEquivalentPlanNode)
+    {
+        return new TableFunctionNode(getSourceLocation(), getId(), statsEquivalentPlanNode, name, arguments, outputVariables, sources, tableArgumentProperties, handle);
+    }
+
+    public static class TableArgumentProperties
+    {
+        private final boolean rowSemantics;
+        private final boolean pruneWhenEmpty;
+        private final boolean passThroughColumns;
+        private final Optional specification;
+
+        @JsonCreator
+        public TableArgumentProperties(
+                @JsonProperty("rowSemantics") boolean rowSemantics,
+                @JsonProperty("pruneWhenEmpty") boolean pruneWhenEmpty,
+                @JsonProperty("passThroughColumns") boolean passThroughColumns,
+                @JsonProperty("specification") Optional specification)
+        {
+            this.rowSemantics = rowSemantics;
+            this.pruneWhenEmpty = pruneWhenEmpty;
+            this.passThroughColumns = passThroughColumns;
+            this.specification = requireNonNull(specification, "specification is null");
+        }
+
+        @JsonProperty
+        public boolean isRowSemantics()
+        {
+            return rowSemantics;
+        }
+
+        @JsonProperty
+        public boolean isPruneWhenEmpty()
+        {
+            return pruneWhenEmpty;
+        }
+
+        @JsonProperty
+        public boolean isPassThroughColumns()
+        {
+            return passThroughColumns;
+        }
+
+        @JsonProperty
+        public Optional specification()
+        {
+            return specification;
+        }
+    }
+}
diff --git a/presto-main-base/src/main/java/com/facebook/presto/sql/planner/planPrinter/PlanPrinter.java b/presto-main-base/src/main/java/com/facebook/presto/sql/planner/planPrinter/PlanPrinter.java
index 8eaacb43ff006..e536f99c5f67f 100644
--- a/presto-main-base/src/main/java/com/facebook/presto/sql/planner/planPrinter/PlanPrinter.java
+++ b/presto-main-base/src/main/java/com/facebook/presto/sql/planner/planPrinter/PlanPrinter.java
@@ -98,6 +98,7 @@
 import com.facebook.presto.sql.planner.plan.SampleNode;
 import com.facebook.presto.sql.planner.plan.SequenceNode;
 import com.facebook.presto.sql.planner.plan.StatisticsWriterNode;
+import com.facebook.presto.sql.planner.plan.TableFunctionNode;
 import com.facebook.presto.sql.planner.plan.TableWriterMergeNode;
 import com.facebook.presto.sql.planner.plan.TopNRowNumberNode;
 import com.facebook.presto.sql.planner.plan.UpdateNode;
@@ -1317,6 +1318,20 @@ public Void visitLateralJoin(LateralJoinNode node, Void context)
         public Void visitOffset(OffsetNode node, Void context)
         {
             addNode(node, "Offset", format("[%s]", node.getCount()));
+            return processChildren(node, context);
+        }
+
+        @Override
+        public Void visitTableFunction(TableFunctionNode node, Void context)
+        {
+            NodeRepresentation nodeOutput = addNode(
+                    node,
+                    "TableFunction",
+                    node.getName());
+
+            checkArgument(
+                    node.getSources().isEmpty() && node.getTableArgumentProperties().isEmpty(),
+                    "Table or descriptor arguments are not yet supported in PlanPrinter");
 
             return processChildren(node, context);
         }
diff --git a/presto-main-base/src/main/java/com/facebook/presto/sql/planner/sanity/ValidateDependenciesChecker.java b/presto-main-base/src/main/java/com/facebook/presto/sql/planner/sanity/ValidateDependenciesChecker.java
index 60f7140716018..333b4c9c4fd76 100644
--- a/presto-main-base/src/main/java/com/facebook/presto/sql/planner/sanity/ValidateDependenciesChecker.java
+++ b/presto-main-base/src/main/java/com/facebook/presto/sql/planner/sanity/ValidateDependenciesChecker.java
@@ -69,6 +69,7 @@
 import com.facebook.presto.sql.planner.plan.SampleNode;
 import com.facebook.presto.sql.planner.plan.SequenceNode;
 import com.facebook.presto.sql.planner.plan.StatisticsWriterNode;
+import com.facebook.presto.sql.planner.plan.TableFunctionNode;
 import com.facebook.presto.sql.planner.plan.TableWriterMergeNode;
 import com.facebook.presto.sql.planner.plan.TopNRowNumberNode;
 import com.facebook.presto.sql.planner.plan.UpdateNode;
@@ -116,6 +117,12 @@ public Void visitPlan(PlanNode node, Set boundVaria
             throw new UnsupportedOperationException("not yet implemented: " + node.getClass().getName());
         }
 
+        @Override
+        public Void visitTableFunction(TableFunctionNode node, Set boundSymbols)
+        {
+            return null;
+        }
+
         @Override
         public Void visitExplainAnalyze(ExplainAnalyzeNode node, Set boundVariables)
         {
diff --git a/presto-main-base/src/main/java/com/facebook/presto/sql/planner/sanity/ValidateStreamingAggregations.java b/presto-main-base/src/main/java/com/facebook/presto/sql/planner/sanity/ValidateStreamingAggregations.java
index cfe040826bdcf..f5ca4ab176918 100644
--- a/presto-main-base/src/main/java/com/facebook/presto/sql/planner/sanity/ValidateStreamingAggregations.java
+++ b/presto-main-base/src/main/java/com/facebook/presto/sql/planner/sanity/ValidateStreamingAggregations.java
@@ -86,8 +86,10 @@ public Void visitAggregation(AggregationNode node, Void context)
 
             List> desiredProperties = ImmutableList.of(new GroupingProperty<>(node.getPreGroupedVariables()));
             Iterator>> matchIterator = LocalProperties.match(properties.getLocalProperties(), desiredProperties).iterator();
+            Iterator>> additionalMatchIterator = LocalProperties.match(properties.getAdditionalLocalProperties(), desiredProperties).iterator();
             Optional> unsatisfiedRequirement = Iterators.getOnlyElement(matchIterator);
-            checkArgument(!unsatisfiedRequirement.isPresent(), "Streaming aggregation with input not grouped on the grouping keys");
+            Optional> additionalUnsatisfiedRequirement = Iterators.getOnlyElement(additionalMatchIterator);
+            checkArgument(!unsatisfiedRequirement.isPresent() || !additionalUnsatisfiedRequirement.isPresent(), "Streaming aggregation with input not grouped on the grouping keys");
             return null;
         }
     }
diff --git a/presto-main-base/src/main/java/com/facebook/presto/sql/relational/Expressions.java b/presto-main-base/src/main/java/com/facebook/presto/sql/relational/Expressions.java
index 3f4340ffdd445..373ce2bfd65a2 100644
--- a/presto-main-base/src/main/java/com/facebook/presto/sql/relational/Expressions.java
+++ b/presto-main-base/src/main/java/com/facebook/presto/sql/relational/Expressions.java
@@ -13,7 +13,6 @@
  */
 package com.facebook.presto.sql.relational;
 
-import com.facebook.presto.common.QualifiedObjectName;
 import com.facebook.presto.common.function.OperatorType;
 import com.facebook.presto.common.type.Type;
 import com.facebook.presto.metadata.CastType;
@@ -154,12 +153,6 @@ public static CallExpression call(FunctionAndTypeManager functionAndTypeManager,
         return call(name, functionHandle, returnType, arguments);
     }
 
-    public static CallExpression call(FunctionAndTypeManager functionAndTypeManager, QualifiedObjectName qualifiedObjectName, Type returnType, List arguments)
-    {
-        FunctionHandle functionHandle = functionAndTypeManager.lookupFunction(qualifiedObjectName, fromTypes(arguments.stream().map(RowExpression::getType).collect(toImmutableList())));
-        return call(String.valueOf(qualifiedObjectName), functionHandle, returnType, arguments);
-    }
-
     public static CallExpression call(FunctionAndTypeResolver functionAndTypeResolver, String name, Type returnType, RowExpression... arguments)
     {
         FunctionHandle functionHandle = functionAndTypeResolver.lookupFunction(name, fromTypes(Arrays.stream(arguments).map(RowExpression::getType).collect(toImmutableList())));
diff --git a/presto-main-base/src/main/java/com/facebook/presto/sql/relational/FunctionResolution.java b/presto-main-base/src/main/java/com/facebook/presto/sql/relational/FunctionResolution.java
index a3fe25d7b65fa..337582257bad4 100644
--- a/presto-main-base/src/main/java/com/facebook/presto/sql/relational/FunctionResolution.java
+++ b/presto-main-base/src/main/java/com/facebook/presto/sql/relational/FunctionResolution.java
@@ -281,7 +281,7 @@ public boolean isEqualsFunction(FunctionHandle functionHandle)
     @Override
     public FunctionHandle subscriptFunction(Type baseType, Type indexType)
     {
-        return functionAndTypeResolver.lookupFunction(SUBSCRIPT.getFunctionName().getObjectName(), fromTypes(baseType, indexType));
+        return functionAndTypeResolver.resolveOperator(SUBSCRIPT, fromTypes(baseType, indexType));
     }
 
     @Override
diff --git a/presto-main-base/src/main/java/com/facebook/presto/sql/rewrite/NativeExecutionTypeRewrite.java b/presto-main-base/src/main/java/com/facebook/presto/sql/rewrite/NativeExecutionTypeRewrite.java
index c0abeb230ae87..d458bccba14e8 100644
--- a/presto-main-base/src/main/java/com/facebook/presto/sql/rewrite/NativeExecutionTypeRewrite.java
+++ b/presto-main-base/src/main/java/com/facebook/presto/sql/rewrite/NativeExecutionTypeRewrite.java
@@ -16,7 +16,6 @@
 import com.facebook.airlift.log.Logger;
 import com.facebook.presto.Session;
 import com.facebook.presto.SystemSessionProperties;
-import com.facebook.presto.UnknownTypeException;
 import com.facebook.presto.common.type.BigintEnumType;
 import com.facebook.presto.common.type.EnumType;
 import com.facebook.presto.common.type.Type;
@@ -25,6 +24,7 @@
 import com.facebook.presto.metadata.Metadata;
 import com.facebook.presto.spi.WarningCollector;
 import com.facebook.presto.spi.security.AccessControl;
+import com.facebook.presto.spi.type.UnknownTypeException;
 import com.facebook.presto.sql.analyzer.FunctionAndTypeResolver;
 import com.facebook.presto.sql.analyzer.QueryExplainer;
 import com.facebook.presto.sql.analyzer.SemanticException;
diff --git a/presto-main-base/src/main/java/com/facebook/presto/tdigest/TDigest.java b/presto-main-base/src/main/java/com/facebook/presto/tdigest/TDigest.java
index 80a750a887fee..0a7977cd6fa67 100644
--- a/presto-main-base/src/main/java/com/facebook/presto/tdigest/TDigest.java
+++ b/presto-main-base/src/main/java/com/facebook/presto/tdigest/TDigest.java
@@ -203,6 +203,13 @@ public static TDigest createTDigest(Slice slice)
             r.mean = new double[r.activeCentroids];
             sliceInput.readBytes(wrappedDoubleArray(r.weight), r.activeCentroids * SIZE_OF_DOUBLE);
             sliceInput.readBytes(wrappedDoubleArray(r.mean), r.activeCentroids * SIZE_OF_DOUBLE);
+
+            // Validate deserialized TDigest data
+            for (int i = 0; i < r.activeCentroids; i++) {
+                checkArgument(!isNaN(r.mean[i]), "Deserialized t-digest contains NaN mean value");
+                checkArgument(r.weight[i] > 0, "weight must be > 0");
+            }
+
             sliceInput.close();
             return r;
         }
@@ -712,6 +719,12 @@ public long estimatedInMemorySizeInBytes()
 
     public Slice serialize()
     {
+        // Validate data before serialization
+        for (int i = 0; i < activeCentroids; i++) {
+            checkArgument(!isNaN(mean[i]), "Cannot serialize t-digest with NaN mean value");
+            checkArgument(weight[i] > 0, "Cannot serialize t-digest with non-positive weight");
+        }
+
         SliceOutput sliceOutput = new DynamicSliceOutput(toIntExact(estimatedSerializedSizeInBytes()));
 
         sliceOutput.writeByte(1); // version 1 of T-Digest serialization
diff --git a/presto-main-base/src/main/java/com/facebook/presto/testing/MaterializedResult.java b/presto-main-base/src/main/java/com/facebook/presto/testing/MaterializedResult.java
index d46e595a537d4..848bc0de0ec0a 100644
--- a/presto-main-base/src/main/java/com/facebook/presto/testing/MaterializedResult.java
+++ b/presto-main-base/src/main/java/com/facebook/presto/testing/MaterializedResult.java
@@ -18,6 +18,7 @@
 import com.facebook.presto.common.PageBuilder;
 import com.facebook.presto.common.block.Block;
 import com.facebook.presto.common.block.BlockBuilder;
+import com.facebook.presto.common.transaction.TransactionId;
 import com.facebook.presto.common.type.ArrayType;
 import com.facebook.presto.common.type.CharType;
 import com.facebook.presto.common.type.MapType;
@@ -99,11 +100,13 @@ public class MaterializedResult
     private final Set resetSessionProperties;
     private final Optional updateType;
     private final OptionalLong updateCount;
+    private final Optional startedTransactionId;
+    private final boolean clearTransactionId;
     private final List warnings;
 
     public MaterializedResult(List rows, List types)
     {
-        this(rows, types, ImmutableMap.of(), ImmutableSet.of(), Optional.empty(), OptionalLong.empty(), ImmutableList.of());
+        this(rows, types, ImmutableMap.of(), ImmutableSet.of(), Optional.empty(), OptionalLong.empty(), Optional.empty(), false, ImmutableList.of());
     }
 
     public MaterializedResult(
@@ -113,6 +116,8 @@ public MaterializedResult(
             Set resetSessionProperties,
             Optional updateType,
             OptionalLong updateCount,
+            Optional startedTransactionId,
+            boolean clearTransactionId,
             List warnings)
     {
         this.rows = ImmutableList.copyOf(requireNonNull(rows, "rows is null"));
@@ -121,6 +126,8 @@ public MaterializedResult(
         this.resetSessionProperties = ImmutableSet.copyOf(requireNonNull(resetSessionProperties, "resetSessionProperties is null"));
         this.updateType = requireNonNull(updateType, "updateType is null");
         this.updateCount = requireNonNull(updateCount, "updateCount is null");
+        this.startedTransactionId = requireNonNull(startedTransactionId, "startedTransactionId is null");
+        this.clearTransactionId = clearTransactionId;
         this.warnings = requireNonNull(warnings, "warnings is null");
     }
 
@@ -165,6 +172,16 @@ public OptionalLong getUpdateCount()
         return updateCount;
     }
 
+    public Optional getStartedTransactionId()
+    {
+        return startedTransactionId;
+    }
+
+    public boolean isClearTransactionId()
+    {
+        return clearTransactionId;
+    }
+
     public List getWarnings()
     {
         return warnings;
@@ -185,13 +202,15 @@ public boolean equals(Object obj)
                 Objects.equals(setSessionProperties, o.setSessionProperties) &&
                 Objects.equals(resetSessionProperties, o.resetSessionProperties) &&
                 Objects.equals(updateType, o.updateType) &&
-                Objects.equals(updateCount, o.updateCount);
+                Objects.equals(updateCount, o.updateCount) &&
+                Objects.equals(startedTransactionId, o.startedTransactionId) &&
+                Objects.equals(clearTransactionId, o.clearTransactionId);
     }
 
     @Override
     public int hashCode()
     {
-        return Objects.hash(rows, types, setSessionProperties, resetSessionProperties, updateType, updateCount);
+        return Objects.hash(rows, types, setSessionProperties, resetSessionProperties, updateType, updateCount, startedTransactionId, clearTransactionId);
     }
 
     @Override
@@ -204,6 +223,8 @@ public String toString()
                 .add("resetSessionProperties", resetSessionProperties)
                 .add("updateType", updateType.orElse(null))
                 .add("updateCount", updateCount.isPresent() ? updateCount.getAsLong() : null)
+                .add("startedTransactionId", startedTransactionId.orElse(null))
+                .add("clearTransactionId", clearTransactionId)
                 .omitNullValues()
                 .toString();
     }
@@ -360,6 +381,8 @@ public MaterializedResult toTestTypes()
                 resetSessionProperties,
                 updateType,
                 updateCount,
+                startedTransactionId,
+                clearTransactionId,
                 warnings);
     }
 
diff --git a/presto-main-base/src/main/java/com/facebook/presto/testing/QueryRunner.java b/presto-main-base/src/main/java/com/facebook/presto/testing/QueryRunner.java
index 5ab9fcbb94a71..f6c1a0b9b1405 100644
--- a/presto-main-base/src/main/java/com/facebook/presto/testing/QueryRunner.java
+++ b/presto-main-base/src/main/java/com/facebook/presto/testing/QueryRunner.java
@@ -126,6 +126,11 @@ default void loadPlanCheckerProviderManager(String planCheckerProviderName, Map<
         throw new UnsupportedOperationException();
     }
 
+    default void triggerConflictCheckWithBuiltInFunctions()
+    {
+        throw new UnsupportedOperationException();
+    }
+
     class MaterializedResultWithPlan
     {
         private final MaterializedResult materializedResult;
diff --git a/presto-main-base/src/main/java/com/facebook/presto/transaction/InMemoryTransactionManager.java b/presto-main-base/src/main/java/com/facebook/presto/transaction/InMemoryTransactionManager.java
index 8afc8c07370e4..bd3a4dbd4f1fd 100644
--- a/presto-main-base/src/main/java/com/facebook/presto/transaction/InMemoryTransactionManager.java
+++ b/presto-main-base/src/main/java/com/facebook/presto/transaction/InMemoryTransactionManager.java
@@ -19,6 +19,7 @@
 import com.facebook.airlift.units.Duration;
 import com.facebook.presto.common.transaction.TransactionId;
 import com.facebook.presto.metadata.Catalog;
+import com.facebook.presto.metadata.Catalog.CatalogContext;
 import com.facebook.presto.metadata.CatalogManager;
 import com.facebook.presto.metadata.CatalogMetadata;
 import com.facebook.presto.spi.ConnectorId;
@@ -214,6 +215,12 @@ public Map getCatalogNames(TransactionId transactionId)
         return getTransactionMetadata(transactionId).getCatalogNames();
     }
 
+    @Override
+    public Map getCatalogNamesWithConnectorContext(TransactionId transactionId)
+    {
+        return getTransactionMetadata(transactionId).getCatalogNamesWithConnectorContext();
+    }
+
     @Override
     public Optional getOptionalCatalogMetadata(TransactionId transactionId, String catalogName)
     {
@@ -462,6 +469,20 @@ private synchronized Map getCatalogNames()
             return ImmutableMap.copyOf(catalogNames);
         }
 
+        private synchronized Map getCatalogNamesWithConnectorContext()
+        {
+            Map catalogNamesWithConnectorContext = new HashMap<>();
+            catalogByName.values().stream()
+                    .filter(Optional::isPresent)
+                    .map(Optional::get)
+                    .forEach(catalog -> catalogNamesWithConnectorContext.put(catalog.getCatalogName(), catalog.getCatalogContext()));
+
+            catalogManager.getCatalogs().stream()
+                    .forEach(catalog -> catalogNamesWithConnectorContext.putIfAbsent(catalog.getCatalogName(), catalog.getCatalogContext()));
+
+            return ImmutableMap.copyOf(catalogNamesWithConnectorContext);
+        }
+
         private synchronized Optional getConnectorId(String catalogName)
         {
             Optional catalog = catalogByName.get(catalogName);
diff --git a/presto-main-base/src/main/java/com/facebook/presto/transaction/TransactionManager.java b/presto-main-base/src/main/java/com/facebook/presto/transaction/TransactionManager.java
index b7d3e0ccdca4a..ecacaf515e0ca 100644
--- a/presto-main-base/src/main/java/com/facebook/presto/transaction/TransactionManager.java
+++ b/presto-main-base/src/main/java/com/facebook/presto/transaction/TransactionManager.java
@@ -15,6 +15,7 @@
 
 import com.facebook.presto.Session;
 import com.facebook.presto.common.transaction.TransactionId;
+import com.facebook.presto.metadata.Catalog.CatalogContext;
 import com.facebook.presto.metadata.CatalogMetadata;
 import com.facebook.presto.spi.ConnectorId;
 import com.facebook.presto.spi.connector.ConnectorTransactionHandle;
@@ -22,6 +23,7 @@
 import com.facebook.presto.spi.function.FunctionNamespaceTransactionHandle;
 import com.facebook.presto.spi.security.AccessControl;
 import com.facebook.presto.spi.transaction.IsolationLevel;
+import com.google.common.collect.ImmutableMap;
 import com.google.common.util.concurrent.ListenableFuture;
 
 import java.util.List;
@@ -47,6 +49,11 @@ public interface TransactionManager
 
     Map getCatalogNames(TransactionId transactionId);
 
+    default Map getCatalogNamesWithConnectorContext(TransactionId transactionId)
+    {
+        return ImmutableMap.of();
+    }
+
     Optional getOptionalCatalogMetadata(TransactionId transactionId, String catalogName);
 
     void enableRollback(TransactionId transactionId);
diff --git a/presto-main-base/src/main/java/com/facebook/presto/type/BigintOperators.java b/presto-main-base/src/main/java/com/facebook/presto/type/BigintOperators.java
index 92dc55dd997db..4689354736795 100644
--- a/presto-main-base/src/main/java/com/facebook/presto/type/BigintOperators.java
+++ b/presto-main-base/src/main/java/com/facebook/presto/type/BigintOperators.java
@@ -20,6 +20,7 @@
 import com.facebook.presto.spi.function.BlockIndex;
 import com.facebook.presto.spi.function.BlockPosition;
 import com.facebook.presto.spi.function.IsNull;
+import com.facebook.presto.spi.function.LiteralParameter;
 import com.facebook.presto.spi.function.LiteralParameters;
 import com.facebook.presto.spi.function.ScalarOperator;
 import com.facebook.presto.spi.function.SqlNullable;
diff --git a/presto-main-base/src/main/java/com/facebook/presto/type/LikeFunctions.java b/presto-main-base/src/main/java/com/facebook/presto/type/LikeFunctions.java
index 7b89abc3bf48d..672d8b3542ff8 100644
--- a/presto-main-base/src/main/java/com/facebook/presto/type/LikeFunctions.java
+++ b/presto-main-base/src/main/java/com/facebook/presto/type/LikeFunctions.java
@@ -16,6 +16,7 @@
 import com.facebook.presto.common.function.OperatorType;
 import com.facebook.presto.common.type.StandardTypes;
 import com.facebook.presto.spi.PrestoException;
+import com.facebook.presto.spi.function.LiteralParameter;
 import com.facebook.presto.spi.function.LiteralParameters;
 import com.facebook.presto.spi.function.ScalarFunction;
 import com.facebook.presto.spi.function.ScalarOperator;
diff --git a/presto-main-base/src/main/java/com/facebook/presto/util/JsonUtil.java b/presto-main-base/src/main/java/com/facebook/presto/util/JsonUtil.java
index 4b48dcbd8fff8..e8e7cb4802de2 100644
--- a/presto-main-base/src/main/java/com/facebook/presto/util/JsonUtil.java
+++ b/presto-main-base/src/main/java/com/facebook/presto/util/JsonUtil.java
@@ -13,6 +13,7 @@
  */
 package com.facebook.presto.util;
 
+import com.facebook.airlift.json.JsonObjectMapperProvider;
 import com.facebook.presto.common.block.Block;
 import com.facebook.presto.common.block.BlockBuilder;
 import com.facebook.presto.common.block.SingleRowBlockWriter;
@@ -41,6 +42,7 @@
 import com.fasterxml.jackson.databind.ObjectMapper;
 import com.google.common.primitives.Shorts;
 import com.google.common.primitives.SignedBytes;
+import io.airlift.slice.DynamicSliceOutput;
 import io.airlift.slice.Slice;
 import io.airlift.slice.SliceOutput;
 import io.airlift.slice.Slices;
@@ -86,6 +88,7 @@
 import static com.fasterxml.jackson.core.JsonToken.FIELD_NAME;
 import static com.fasterxml.jackson.core.JsonToken.START_ARRAY;
 import static com.fasterxml.jackson.core.JsonToken.START_OBJECT;
+import static com.fasterxml.jackson.databind.SerializationFeature.ORDER_MAP_ENTRIES_BY_KEYS;
 import static com.google.common.base.Preconditions.checkArgument;
 import static com.google.common.base.Verify.verify;
 import static it.unimi.dsi.fastutil.HashCommon.arraySize;
@@ -106,6 +109,7 @@ public final class JsonUtil
     // `OBJECT_MAPPER.writeValueAsString(parser.readValueAsTree());` preserves input order.
     // Be aware. Using it arbitrarily can produce invalid json (ordered by key is required in Presto).
     private static final ObjectMapper OBJECT_MAPPED_UNORDERED = new ObjectMapper(JSON_FACTORY);
+    private static final ObjectMapper OBJECT_MAPPED_SORTED = new JsonObjectMapperProvider().get().configure(ORDER_MAP_ENTRIES_BY_KEYS, true);
 
     private static final int MAX_JSON_LENGTH_IN_ERROR_MESSAGE = 10_000;
 
@@ -956,8 +960,18 @@ static BlockBuilderAppender createBlockBuilderAppender(Type type)
                     return new VarcharBlockBuilderAppender(type);
                 case StandardTypes.JSON:
                     return (parser, blockBuilder, sqlFunctionProperties) -> {
-                        String json = OBJECT_MAPPED_UNORDERED.writeValueAsString(parser.readValueAsTree());
-                        JSON.writeSlice(blockBuilder, Slices.utf8Slice(json));
+                        Slice slice = Slices.utf8Slice(OBJECT_MAPPED_UNORDERED.writeValueAsString(parser.readValueAsTree()));
+                        try (JsonParser jsonParser = createJsonParser(JSON_FACTORY, slice)) {
+                            SliceOutput dynamicSliceOutput = new DynamicSliceOutput(slice.length());
+                            OBJECT_MAPPED_SORTED.writeValue((OutputStream) dynamicSliceOutput, OBJECT_MAPPED_SORTED.readValue(jsonParser, Object.class));
+                            // nextToken() returns null if the input is parsed correctly,
+                            // but will throw an exception if there are trailing characters.
+                            jsonParser.nextToken();
+                            JSON.writeSlice(blockBuilder, dynamicSliceOutput.slice());
+                        }
+                        catch (Exception e) {
+                            throw new PrestoException(INVALID_FUNCTION_ARGUMENT, format("Cannot convert '%s' to JSON", slice.toStringUtf8()));
+                        }
                     };
                 case StandardTypes.ARRAY:
                     return new ArrayBlockBuilderAppender(createBlockBuilderAppender(((ArrayType) type).getElementType()));
diff --git a/presto-main-base/src/test/java/com/facebook/presto/connector/tvf/TestTVFConnectorColumnHandle.java b/presto-main-base/src/test/java/com/facebook/presto/connector/tvf/TestTVFConnectorColumnHandle.java
new file mode 100644
index 0000000000000..063d4e37a839c
--- /dev/null
+++ b/presto-main-base/src/test/java/com/facebook/presto/connector/tvf/TestTVFConnectorColumnHandle.java
@@ -0,0 +1,81 @@
+/*
+ * Licensed 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 com.facebook.presto.connector.tvf;
+
+import com.facebook.presto.common.type.Type;
+import com.facebook.presto.spi.ColumnHandle;
+import com.fasterxml.jackson.annotation.JsonCreator;
+import com.fasterxml.jackson.annotation.JsonProperty;
+
+import java.util.Objects;
+
+import static com.google.common.base.MoreObjects.toStringHelper;
+import static java.util.Objects.requireNonNull;
+
+public class TestTVFConnectorColumnHandle
+        implements ColumnHandle
+{
+    private final String name;
+    private final Type type;
+
+    @JsonCreator
+    public TestTVFConnectorColumnHandle(
+            @JsonProperty("name") String name,
+            @JsonProperty("type") Type type)
+    {
+        this.name = requireNonNull(name, "name is null");
+        this.type = requireNonNull(type, "type is null");
+    }
+
+    @JsonProperty
+    public String getName()
+    {
+        return name;
+    }
+
+    @JsonProperty
+    public Type getType()
+    {
+        return type;
+    }
+
+    @Override
+    public String toString()
+    {
+        return toStringHelper(this)
+                .add("name", name)
+                .add("type", type)
+                .toString();
+    }
+
+    @Override
+    public boolean equals(Object o)
+    {
+        if (this == o) {
+            return true;
+        }
+        if ((o == null) || (getClass() != o.getClass())) {
+            return false;
+        }
+        TestTVFConnectorColumnHandle other = (TestTVFConnectorColumnHandle) o;
+        return Objects.equals(name, other.name) &&
+                Objects.equals(type, other.type);
+    }
+
+    @Override
+    public int hashCode()
+    {
+        return Objects.hash(name, type);
+    }
+}
diff --git a/presto-main-base/src/test/java/com/facebook/presto/connector/tvf/TestTVFConnectorFactory.java b/presto-main-base/src/test/java/com/facebook/presto/connector/tvf/TestTVFConnectorFactory.java
new file mode 100644
index 0000000000000..2674c87d28cc6
--- /dev/null
+++ b/presto-main-base/src/test/java/com/facebook/presto/connector/tvf/TestTVFConnectorFactory.java
@@ -0,0 +1,452 @@
+/*
+ * Licensed 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 com.facebook.presto.connector.tvf;
+
+import com.facebook.presto.common.RuntimeStats;
+import com.facebook.presto.common.predicate.TupleDomain;
+import com.facebook.presto.common.type.Type;
+import com.facebook.presto.spi.ColumnHandle;
+import com.facebook.presto.spi.ColumnMetadata;
+import com.facebook.presto.spi.ConnectorHandleResolver;
+import com.facebook.presto.spi.ConnectorPageSource;
+import com.facebook.presto.spi.ConnectorSession;
+import com.facebook.presto.spi.ConnectorSplit;
+import com.facebook.presto.spi.ConnectorSplitSource;
+import com.facebook.presto.spi.ConnectorTableHandle;
+import com.facebook.presto.spi.ConnectorTableLayout;
+import com.facebook.presto.spi.ConnectorTableLayoutHandle;
+import com.facebook.presto.spi.ConnectorTableLayoutResult;
+import com.facebook.presto.spi.ConnectorTableMetadata;
+import com.facebook.presto.spi.ConnectorViewDefinition;
+import com.facebook.presto.spi.Constraint;
+import com.facebook.presto.spi.FixedSplitSource;
+import com.facebook.presto.spi.HostAddress;
+import com.facebook.presto.spi.InMemoryRecordSet;
+import com.facebook.presto.spi.NodeProvider;
+import com.facebook.presto.spi.RecordPageSource;
+import com.facebook.presto.spi.SchemaTableName;
+import com.facebook.presto.spi.SchemaTablePrefix;
+import com.facebook.presto.spi.SplitContext;
+import com.facebook.presto.spi.connector.Connector;
+import com.facebook.presto.spi.connector.ConnectorContext;
+import com.facebook.presto.spi.connector.ConnectorFactory;
+import com.facebook.presto.spi.connector.ConnectorMetadata;
+import com.facebook.presto.spi.connector.ConnectorPageSourceProvider;
+import com.facebook.presto.spi.connector.ConnectorRecordSetProvider;
+import com.facebook.presto.spi.connector.ConnectorSplitManager;
+import com.facebook.presto.spi.connector.ConnectorTransactionHandle;
+import com.facebook.presto.spi.connector.TableFunctionApplicationResult;
+import com.facebook.presto.spi.function.table.ConnectorTableFunction;
+import com.facebook.presto.spi.function.table.ConnectorTableFunctionHandle;
+import com.facebook.presto.spi.schedule.NodeSelectionStrategy;
+import com.facebook.presto.spi.statistics.TableStatistics;
+import com.facebook.presto.spi.transaction.IsolationLevel;
+import com.facebook.presto.tpch.TpchColumnHandle;
+import com.facebook.presto.tpch.TpchRecordSetProvider;
+import com.google.common.collect.ImmutableList;
+import com.google.common.collect.ImmutableMap;
+import com.google.common.collect.ImmutableSet;
+
+import java.util.Collections;
+import java.util.List;
+import java.util.Map;
+import java.util.Optional;
+import java.util.Set;
+import java.util.function.BiFunction;
+import java.util.function.Function;
+import java.util.function.Supplier;
+import java.util.stream.IntStream;
+
+import static com.facebook.presto.common.type.VarcharType.createUnboundedVarcharType;
+import static com.facebook.presto.spi.schedule.NodeSelectionStrategy.NO_PREFERENCE;
+import static com.google.common.collect.ImmutableList.toImmutableList;
+import static com.google.common.collect.ImmutableMap.toImmutableMap;
+import static java.util.Objects.requireNonNull;
+
+public class TestTVFConnectorFactory
+        implements ConnectorFactory
+{
+    private final Function> listSchemaNames;
+    private final BiFunction> listTables;
+    private final BiFunction> getViews;
+    private final BiFunction> getColumnHandles;
+    private final Supplier getTableStatistics;
+    private final ApplyTableFunction applyTableFunction;
+    private final Set tableFunctions;
+
+    private TestTVFConnectorFactory(
+            Function> listSchemaNames,
+            BiFunction> listTables,
+            BiFunction> getViews,
+            BiFunction> getColumnHandles,
+            Supplier getTableStatistics,
+            ApplyTableFunction applyTableFunction,
+            Set tableFunctions)
+    {
+        this.listSchemaNames = requireNonNull(listSchemaNames, "listSchemaNames is null");
+        this.listTables = requireNonNull(listTables, "listTables is null");
+        this.getViews = requireNonNull(getViews, "getViews is null");
+        this.getColumnHandles = requireNonNull(getColumnHandles, "getColumnHandles is null");
+        this.getTableStatistics = requireNonNull(getTableStatistics, "getTableStatistics is null");
+        this.applyTableFunction = requireNonNull(applyTableFunction, "applyTableFunction is null");
+        this.tableFunctions = requireNonNull(tableFunctions, "tableFunctions is null");
+    }
+
+    @Override
+    public String getName()
+    {
+        return "testTVF";
+    }
+
+    @Override
+    public ConnectorHandleResolver getHandleResolver()
+    {
+        return new TestTVFHandleResolver();
+    }
+
+    @Override
+    public Connector create(String catalogName, Map config, ConnectorContext context)
+    {
+        return new TestTVFConnector(context, listSchemaNames, listTables, getViews, getColumnHandles, getTableStatistics, applyTableFunction, tableFunctions);
+    }
+
+    public static Builder builder()
+    {
+        return new Builder();
+    }
+
+    public static Function> defaultGetColumns()
+    {
+        return table -> IntStream.range(0, 100)
+                .boxed()
+                .map(i -> ColumnMetadata.builder().setName("column_" + i).setType(createUnboundedVarcharType()).build())
+                .collect(toImmutableList());
+    }
+
+    @FunctionalInterface
+    public interface ApplyTableFunction
+    {
+        Optional> apply(ConnectorSession session, ConnectorTableFunctionHandle handle);
+    }
+
+    public static class TestTVFConnector
+            implements Connector
+    {
+        private static final String DELETE_ROW_ID = "delete_row_id";
+        private static final String UPDATE_ROW_ID = "update_row_id";
+        private static final String MERGE_ROW_ID = "merge_row_id";
+
+        private final ConnectorContext context;
+        private final Function> listSchemaNames;
+        private final BiFunction> listTables;
+        private final BiFunction> getViews;
+        private final BiFunction> getColumnHandles;
+        private final Supplier getTableStatistics;
+        private final ApplyTableFunction applyTableFunction;
+        private final Set tableFunctions;
+
+        public TestTVFConnector(
+                ConnectorContext context,
+                Function> listSchemaNames,
+                BiFunction> listTables,
+                BiFunction> getViews,
+                BiFunction> getColumnHandles,
+                Supplier getTableStatistics,
+                ApplyTableFunction applyTableFunction,
+                Set tableFunctions)
+        {
+            this.context = requireNonNull(context, "context is null");
+            this.listSchemaNames = requireNonNull(listSchemaNames, "listSchemaNames is null");
+            this.listTables = requireNonNull(listTables, "listTables is null");
+            this.getViews = requireNonNull(getViews, "getViews is null");
+            this.getColumnHandles = requireNonNull(getColumnHandles, "getColumnHandles is null");
+            this.getTableStatistics = requireNonNull(getTableStatistics, "getTableStatistics is null");
+            this.applyTableFunction = requireNonNull(applyTableFunction, "applyTableFunction is null");
+            this.tableFunctions = requireNonNull(tableFunctions, "tableFunctions is null");
+        }
+
+        @Override
+        public ConnectorTransactionHandle beginTransaction(IsolationLevel isolationLevel, boolean readOnly)
+        {
+            return TestTVFConnectorTransactionHandle.INSTANCE;
+        }
+
+        @Override
+        public ConnectorMetadata getMetadata(ConnectorTransactionHandle transaction)
+        {
+            return new TestTVFConnectorMetadata();
+        }
+
+        public enum TestTVFConnectorSplit
+                implements ConnectorSplit
+        {
+            TEST_TVF_CONNECTOR_SPLIT;
+
+            @Override
+            public NodeSelectionStrategy getNodeSelectionStrategy()
+            {
+                return NO_PREFERENCE;
+            }
+
+            @Override
+            public List getPreferredNodes(NodeProvider nodeProvider)
+            {
+                return Collections.emptyList();
+            }
+
+            @Override
+            public Object getInfo()
+            {
+                return null;
+            }
+        }
+
+        @Override
+        public ConnectorSplitManager getSplitManager()
+        {
+            return new ConnectorSplitManager()
+            {
+                @Override
+                public ConnectorSplitSource getSplits(ConnectorTransactionHandle transactionHandle, ConnectorSession session, ConnectorTableLayoutHandle layout, SplitSchedulingContext splitSchedulingContext)
+                {
+                    return new FixedSplitSource(Collections.singleton(TestTVFConnectorSplit.TEST_TVF_CONNECTOR_SPLIT));
+                }
+            };
+        }
+
+        @Override
+        public ConnectorRecordSetProvider getRecordSetProvider()
+        {
+            return new TpchRecordSetProvider();
+        }
+
+        @Override
+        public ConnectorPageSourceProvider getPageSourceProvider()
+        {
+            return new TestTVFConnectorPageSourceProvider();
+        }
+
+        @Override
+        public Set getTableFunctions()
+        {
+            return tableFunctions;
+        }
+
+        private class TestTVFConnectorMetadata
+                implements ConnectorMetadata
+        {
+            @Override
+            public List listSchemaNames(ConnectorSession session)
+            {
+                return listSchemaNames.apply(session);
+            }
+
+            @Override
+            public ConnectorTableHandle getTableHandle(ConnectorSession session, SchemaTableName tableName)
+            {
+                return new ConnectorTableHandle() {};
+            }
+
+            @Override
+            public ConnectorTableMetadata getTableMetadata(ConnectorSession session, ConnectorTableHandle tableHandle)
+            {
+                TestTVFConnectorTableHandle table = (TestTVFConnectorTableHandle) tableHandle;
+                return new ConnectorTableMetadata(
+                        table.getTableName(),
+                        defaultGetColumns().apply(table.getTableName()),
+                        ImmutableMap.of());
+            }
+
+            @Override
+            public List listTables(ConnectorSession session, String schemaNameOrNull)
+            {
+                return listTables.apply(session, schemaNameOrNull);
+            }
+
+            public void setTableProperties(ConnectorSession session, ConnectorTableHandle tableHandle, Map properties)
+            {
+            }
+
+            @Override
+            public Map getColumnHandles(ConnectorSession session, ConnectorTableHandle tableHandle)
+            {
+                return (Map) (Map) getColumnHandles.apply(session, tableHandle);
+            }
+
+            @Override
+            public ColumnMetadata getColumnMetadata(ConnectorSession session, ConnectorTableHandle tableHandle, ColumnHandle columnHandle)
+            {
+                if (columnHandle instanceof TestTVFConnectorColumnHandle) {
+                    TestTVFConnectorColumnHandle testTVFColumnHandle = (TestTVFConnectorColumnHandle) columnHandle;
+                    return ColumnMetadata.builder().setName(testTVFColumnHandle.getName()).setType(testTVFColumnHandle.getType()).build();
+                }
+                else {
+                    TpchColumnHandle tpchColumnHandle = (TpchColumnHandle) columnHandle;
+                    return ColumnMetadata.builder().setName(tpchColumnHandle.getColumnName()).setType(tpchColumnHandle.getType()).build();
+                }
+            }
+
+            @Override
+            public Map> listTableColumns(ConnectorSession session, SchemaTablePrefix prefix)
+            {
+                return listTables(session, prefix.getSchemaName()).stream()
+                        .collect(toImmutableMap(table -> table, table -> IntStream.range(0, 100)
+                                .boxed()
+                                .map(i -> ColumnMetadata.builder().setName("column_" + i).setType(createUnboundedVarcharType()).build())
+                                .collect(toImmutableList())));
+            }
+
+            @Override
+            public ConnectorTableLayoutResult getTableLayoutForConstraint(
+                    ConnectorSession session,
+                    ConnectorTableHandle table,
+                    Constraint constraint,
+                    Optional> desiredColumns)
+            {
+                // TODO: Currently not supporting constraints
+                TestTVFTableLayoutHandle tvfLayout = new TestTVFTableLayoutHandle((TestTVFConnectorTableHandle) table, TupleDomain.none());
+                return new ConnectorTableLayoutResult(new ConnectorTableLayout(tvfLayout,
+                        Optional.empty(),
+                        tvfLayout.getPredicate(),
+                        Optional.empty(),
+                        Optional.empty(),
+                        Optional.empty(),
+                        Collections.emptyList(),
+                        Optional.empty()), TupleDomain.none());
+            }
+
+            @Override
+            public ConnectorTableLayout getTableLayout(ConnectorSession session, ConnectorTableLayoutHandle handle)
+            {
+                TestTVFTableLayoutHandle tvfTableLayout = (TestTVFTableLayoutHandle) handle;
+                return new ConnectorTableLayout(tvfTableLayout);
+            }
+
+            @Override
+            public Map getViews(ConnectorSession session, SchemaTablePrefix prefix)
+            {
+                return getViews.apply(session, prefix);
+            }
+
+            @Override
+            public TableStatistics getTableStatistics(ConnectorSession session, ConnectorTableHandle tableHandle, Optional tableLayoutHandle, List columnHandles, Constraint constraint)
+            {
+                return getTableStatistics.get();
+            }
+
+            @Override
+            public Optional> applyTableFunction(ConnectorSession session, ConnectorTableFunctionHandle handle)
+            {
+                return applyTableFunction.apply(session, handle);
+            }
+        }
+
+        private class TestTVFConnectorPageSourceProvider
+                implements ConnectorPageSourceProvider
+        {
+            @Override
+            public ConnectorPageSource createPageSource(ConnectorTransactionHandle transactionHandle, ConnectorSession session, ConnectorSplit split, ConnectorTableLayoutHandle layout, List columns, SplitContext splitContext, RuntimeStats runtimeStats)
+            {
+                TestTVFConnectorTableHandle handle = ((TestTVFTableLayoutHandle) layout).getTable();
+                SchemaTableName tableName = handle.getTableName();
+                List projection = columns.stream()
+                        .map(TestTVFConnectorColumnHandle.class::cast)
+                        .collect(toImmutableList());
+                List types = columns.stream()
+                        .map(TestTVFConnectorColumnHandle.class::cast)
+                        .map(TestTVFConnectorColumnHandle::getType)
+                        .collect(toImmutableList());
+                return new TestTVFConnectorPageSource(new RecordPageSource(new InMemoryRecordSet(types, ImmutableList.of())));
+            }
+
+            private Map getColumnIndexes(SchemaTableName tableName)
+            {
+                ImmutableMap.Builder columnIndexes = ImmutableMap.builder();
+                List columnMetadata = defaultGetColumns().apply(tableName);
+                for (int index = 0; index < columnMetadata.size(); index++) {
+                    columnIndexes.put(columnMetadata.get(index).getName(), index);
+                }
+                return columnIndexes.buildOrThrow();
+            }
+        }
+    }
+
+    public static final class Builder
+    {
+        private Function> listSchemaNames = (session) -> ImmutableList.of();
+        private BiFunction> listTables = (session, schemaName) -> ImmutableList.of();
+        private BiFunction> getViews = (session, schemaTablePrefix) -> ImmutableMap.of();
+        private BiFunction> getColumnHandles = (session, tableHandle) -> {
+            TestTVFConnectorTableHandle table = (TestTVFConnectorTableHandle) tableHandle;
+            return defaultGetColumns().apply(table.getTableName()).stream()
+                    .collect(toImmutableMap(ColumnMetadata::getName, column ->
+                            new TestTVFConnectorColumnHandle(column.getName(), column.getType())));
+        };
+        private Supplier getTableStatistics = TableStatistics::empty;
+        private ApplyTableFunction applyTableFunction = (session, handle) -> Optional.empty();
+        private Set tableFunctions = ImmutableSet.of();
+
+        public Builder withListSchemaNames(Function> listSchemaNames)
+        {
+            this.listSchemaNames = requireNonNull(listSchemaNames, "listSchemaNames is null");
+            return this;
+        }
+
+        public Builder withListTables(BiFunction> listTables)
+        {
+            this.listTables = requireNonNull(listTables, "listTables is null");
+            return this;
+        }
+
+        public Builder withGetViews(BiFunction> getViews)
+        {
+            this.getViews = requireNonNull(getViews, "getViews is null");
+            return this;
+        }
+
+        public Builder withGetColumnHandles(BiFunction> getColumnHandles)
+        {
+            this.getColumnHandles = requireNonNull(getColumnHandles, "getColumnHandles is null");
+            return this;
+        }
+
+        public Builder withGetTableStatistics(Supplier getTableStatistics)
+        {
+            this.getTableStatistics = requireNonNull(getTableStatistics, "getTableStatistics is null");
+            return this;
+        }
+
+        public Builder withApplyTableFunction(ApplyTableFunction applyTableFunction)
+        {
+            this.applyTableFunction = applyTableFunction;
+            return this;
+        }
+
+        public Builder withTableFunctions(Iterable tableFunctions)
+        {
+            this.tableFunctions = ImmutableSet.copyOf(tableFunctions);
+            return this;
+        }
+
+        public TestTVFConnectorFactory build()
+        {
+            return new TestTVFConnectorFactory(listSchemaNames, listTables, getViews, getColumnHandles, getTableStatistics, applyTableFunction, tableFunctions);
+        }
+
+        private static  T notSupported()
+        {
+            throw new UnsupportedOperationException();
+        }
+    }
+}
diff --git a/presto-main-base/src/test/java/com/facebook/presto/connector/tvf/TestTVFConnectorPageSource.java b/presto-main-base/src/test/java/com/facebook/presto/connector/tvf/TestTVFConnectorPageSource.java
new file mode 100644
index 0000000000000..126b23a48e796
--- /dev/null
+++ b/presto-main-base/src/test/java/com/facebook/presto/connector/tvf/TestTVFConnectorPageSource.java
@@ -0,0 +1,82 @@
+/*
+ * Licensed 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 com.facebook.presto.connector.tvf;
+
+import com.facebook.presto.common.Page;
+import com.facebook.presto.spi.ConnectorPageSource;
+
+import java.io.IOException;
+import java.util.concurrent.CompletableFuture;
+
+import static java.util.Objects.requireNonNull;
+
+public class TestTVFConnectorPageSource
+        implements ConnectorPageSource
+{
+    private final ConnectorPageSource delegate;
+
+    public TestTVFConnectorPageSource(ConnectorPageSource delegate)
+    {
+        this.delegate = requireNonNull(delegate, "delegate is null");
+    }
+
+    @Override
+    public long getCompletedBytes()
+    {
+        return delegate.getCompletedBytes();
+    }
+
+    @Override
+    public long getCompletedPositions()
+    {
+        return 0;
+    }
+
+    @Override
+    public long getReadTimeNanos()
+    {
+        return delegate.getReadTimeNanos();
+    }
+
+    @Override
+    public boolean isFinished()
+    {
+        return delegate.isFinished();
+    }
+
+    @Override
+    public Page getNextPage()
+    {
+        return delegate.getNextPage();
+    }
+
+    @Override
+    public long getSystemMemoryUsage()
+    {
+        return 0;
+    }
+
+    @Override
+    public void close()
+            throws IOException
+    {
+        delegate.close();
+    }
+
+    @Override
+    public CompletableFuture isBlocked()
+    {
+        return delegate.isBlocked();
+    }
+}
diff --git a/presto-main-base/src/test/java/com/facebook/presto/connector/tvf/TestTVFConnectorPlugin.java b/presto-main-base/src/test/java/com/facebook/presto/connector/tvf/TestTVFConnectorPlugin.java
new file mode 100644
index 0000000000000..a26ef9c34466a
--- /dev/null
+++ b/presto-main-base/src/test/java/com/facebook/presto/connector/tvf/TestTVFConnectorPlugin.java
@@ -0,0 +1,37 @@
+/*
+ * Licensed 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 com.facebook.presto.connector.tvf;
+
+import com.facebook.presto.spi.Plugin;
+import com.facebook.presto.spi.connector.ConnectorFactory;
+import com.google.common.collect.ImmutableList;
+
+import static java.util.Objects.requireNonNull;
+
+public class TestTVFConnectorPlugin
+        implements Plugin
+{
+    private final ConnectorFactory connectorFactory;
+
+    public TestTVFConnectorPlugin(ConnectorFactory connectorFactory)
+    {
+        this.connectorFactory = requireNonNull(connectorFactory, "connectorFactory is null");
+    }
+
+    @Override
+    public Iterable getConnectorFactories()
+    {
+        return ImmutableList.of(connectorFactory);
+    }
+}
diff --git a/presto-main-base/src/test/java/com/facebook/presto/connector/tvf/TestTVFConnectorTableHandle.java b/presto-main-base/src/test/java/com/facebook/presto/connector/tvf/TestTVFConnectorTableHandle.java
new file mode 100644
index 0000000000000..9c90975cbc82b
--- /dev/null
+++ b/presto-main-base/src/test/java/com/facebook/presto/connector/tvf/TestTVFConnectorTableHandle.java
@@ -0,0 +1,101 @@
+/*
+ * Licensed 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 com.facebook.presto.connector.tvf;
+
+import com.facebook.presto.common.predicate.TupleDomain;
+import com.facebook.presto.spi.ColumnHandle;
+import com.facebook.presto.spi.ConnectorTableHandle;
+import com.facebook.presto.spi.SchemaTableName;
+import com.fasterxml.jackson.annotation.JsonCreator;
+import com.fasterxml.jackson.annotation.JsonProperty;
+import com.google.common.collect.ImmutableList;
+
+import java.util.List;
+import java.util.Objects;
+import java.util.Optional;
+
+import static java.util.Objects.requireNonNull;
+
+public class TestTVFConnectorTableHandle
+        implements ConnectorTableHandle
+{
+    // These are example fields for a Connector's Table Handle.
+    // For other examples, see TpchTableHandle or any other implementations
+    // of ConnectorTableHandle.
+    private final SchemaTableName tableName;
+    private final Optional> columns;
+    private final TupleDomain constraint;
+
+    public TestTVFConnectorTableHandle(SchemaTableName tableName)
+    {
+        this(tableName, Optional.empty(), TupleDomain.all());
+    }
+
+    @JsonCreator
+    public TestTVFConnectorTableHandle(
+            @JsonProperty SchemaTableName tableName,
+            @JsonProperty("columns") Optional> columns,
+            @JsonProperty("constraint") TupleDomain constraint)
+    {
+        this.tableName = requireNonNull(tableName, "tableName is null");
+        requireNonNull(columns, "columns is null");
+        this.columns = columns.map(ImmutableList::copyOf);
+        this.constraint = requireNonNull(constraint, "constraint is null");
+    }
+
+    @JsonProperty
+    public SchemaTableName getTableName()
+    {
+        return tableName;
+    }
+
+    @JsonProperty
+    public Optional> getColumns()
+    {
+        return columns;
+    }
+
+    @JsonProperty
+    public TupleDomain getConstraint()
+    {
+        return constraint;
+    }
+
+    @Override
+    public boolean equals(Object o)
+    {
+        if (this == o) {
+            return true;
+        }
+        if (o == null || getClass() != o.getClass()) {
+            return false;
+        }
+        TestTVFConnectorTableHandle other = (TestTVFConnectorTableHandle) o;
+        return Objects.equals(tableName, other.tableName) &&
+                Objects.equals(constraint, other.constraint) &&
+                Objects.equals(columns, other.columns);
+    }
+
+    @Override
+    public int hashCode()
+    {
+        return Objects.hash(tableName, constraint, columns);
+    }
+
+    @Override
+    public String toString()
+    {
+        return tableName.toString();
+    }
+}
diff --git a/presto-main-base/src/test/java/com/facebook/presto/connector/tvf/TestTVFConnectorTransactionHandle.java b/presto-main-base/src/test/java/com/facebook/presto/connector/tvf/TestTVFConnectorTransactionHandle.java
new file mode 100644
index 0000000000000..668d3fb6575b2
--- /dev/null
+++ b/presto-main-base/src/test/java/com/facebook/presto/connector/tvf/TestTVFConnectorTransactionHandle.java
@@ -0,0 +1,22 @@
+/*
+ * Licensed 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 com.facebook.presto.connector.tvf;
+
+import com.facebook.presto.spi.connector.ConnectorTransactionHandle;
+
+public enum TestTVFConnectorTransactionHandle
+        implements ConnectorTransactionHandle
+{
+    INSTANCE
+}
diff --git a/presto-main-base/src/test/java/com/facebook/presto/connector/tvf/TestTVFHandleResolver.java b/presto-main-base/src/test/java/com/facebook/presto/connector/tvf/TestTVFHandleResolver.java
new file mode 100644
index 0000000000000..1d6771ef0bfb7
--- /dev/null
+++ b/presto-main-base/src/test/java/com/facebook/presto/connector/tvf/TestTVFHandleResolver.java
@@ -0,0 +1,62 @@
+/*
+ * Licensed 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 com.facebook.presto.connector.tvf;
+
+import com.facebook.presto.spi.ColumnHandle;
+import com.facebook.presto.spi.ConnectorHandleResolver;
+import com.facebook.presto.spi.ConnectorSplit;
+import com.facebook.presto.spi.ConnectorTableHandle;
+import com.facebook.presto.spi.ConnectorTableLayoutHandle;
+import com.facebook.presto.spi.connector.ConnectorPartitioningHandle;
+import com.facebook.presto.spi.connector.ConnectorTransactionHandle;
+
+public class TestTVFHandleResolver
+        implements ConnectorHandleResolver
+{
+    @Override
+    public Class getTableHandleClass()
+    {
+        return TestTVFConnectorTableHandle.class;
+    }
+
+    @Override
+    public Class getColumnHandleClass()
+    {
+        return TestTVFConnectorColumnHandle.class;
+    }
+
+    @Override
+    public Class getSplitClass()
+    {
+        return TestTVFConnectorFactory.TestTVFConnector.TestTVFConnectorSplit.class;
+    }
+
+    @Override
+    public Class getTableLayoutHandleClass()
+    {
+        return TestTVFTableLayoutHandle.class;
+    }
+
+    @Override
+    public Class getTransactionHandleClass()
+    {
+        return TestTVFConnectorTransactionHandle.class;
+    }
+
+    @Override
+    public Class getPartitioningHandleClass()
+    {
+        return TestTVFPartitioningHandle.class;
+    }
+}
diff --git a/presto-main-base/src/test/java/com/facebook/presto/connector/tvf/TestTVFPartitioningHandle.java b/presto-main-base/src/test/java/com/facebook/presto/connector/tvf/TestTVFPartitioningHandle.java
new file mode 100644
index 0000000000000..cca6fc0a48bb4
--- /dev/null
+++ b/presto-main-base/src/test/java/com/facebook/presto/connector/tvf/TestTVFPartitioningHandle.java
@@ -0,0 +1,77 @@
+/*
+ * Licensed 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 com.facebook.presto.connector.tvf;
+
+import com.facebook.presto.spi.connector.ConnectorPartitioningHandle;
+import com.fasterxml.jackson.annotation.JsonCreator;
+import com.fasterxml.jackson.annotation.JsonProperty;
+
+import java.util.Objects;
+
+import static com.google.common.base.Preconditions.checkArgument;
+import static java.util.Objects.requireNonNull;
+
+public class TestTVFPartitioningHandle
+        implements ConnectorPartitioningHandle
+{
+    private final String table;
+    private final long totalRows;
+
+    @JsonCreator
+    public TestTVFPartitioningHandle(@JsonProperty("table") String table, @JsonProperty("totalRows") long totalRows)
+    {
+        this.table = requireNonNull(table, "table is null");
+
+        checkArgument(totalRows > 0, "totalRows must be at least 1");
+        this.totalRows = totalRows;
+    }
+
+    @JsonProperty
+    public String getTable()
+    {
+        return table;
+    }
+
+    @JsonProperty
+    public long getTotalRows()
+    {
+        return totalRows;
+    }
+
+    @Override
+    public boolean equals(Object o)
+    {
+        if (this == o) {
+            return true;
+        }
+        if (o == null || getClass() != o.getClass()) {
+            return false;
+        }
+        TestTVFPartitioningHandle that = (TestTVFPartitioningHandle) o;
+        return Objects.equals(table, that.table) &&
+                totalRows == that.totalRows;
+    }
+
+    @Override
+    public int hashCode()
+    {
+        return Objects.hash(table, totalRows);
+    }
+
+    @Override
+    public String toString()
+    {
+        return table + ":" + totalRows;
+    }
+}
diff --git a/presto-main-base/src/test/java/com/facebook/presto/connector/tvf/TestTVFTableLayoutHandle.java b/presto-main-base/src/test/java/com/facebook/presto/connector/tvf/TestTVFTableLayoutHandle.java
new file mode 100644
index 0000000000000..c9f61203cf1de
--- /dev/null
+++ b/presto-main-base/src/test/java/com/facebook/presto/connector/tvf/TestTVFTableLayoutHandle.java
@@ -0,0 +1,66 @@
+/*
+ * Licensed 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 com.facebook.presto.connector.tvf;
+
+import com.facebook.presto.common.plan.PlanCanonicalizationStrategy;
+import com.facebook.presto.common.predicate.TupleDomain;
+import com.facebook.presto.spi.ColumnHandle;
+import com.facebook.presto.spi.ConnectorSplit;
+import com.facebook.presto.spi.ConnectorTableLayoutHandle;
+import com.fasterxml.jackson.annotation.JsonCreator;
+import com.fasterxml.jackson.annotation.JsonProperty;
+import com.google.common.collect.ImmutableMap;
+
+import java.util.Optional;
+
+public class TestTVFTableLayoutHandle
+        implements ConnectorTableLayoutHandle
+{
+    private final TestTVFConnectorTableHandle table;
+    private final TupleDomain predicate;
+
+    @JsonCreator
+    public TestTVFTableLayoutHandle(@JsonProperty("table") TestTVFConnectorTableHandle table, @JsonProperty("predicate") TupleDomain predicate)
+    {
+        this.table = table;
+        this.predicate = predicate;
+    }
+
+    @JsonProperty
+    public TestTVFConnectorTableHandle getTable()
+    {
+        return table;
+    }
+
+    @JsonProperty
+    public TupleDomain getPredicate()
+    {
+        return predicate;
+    }
+
+    @Override
+    public String toString()
+    {
+        return table.toString();
+    }
+
+    @Override
+    public Object getIdentifier(Optional split, PlanCanonicalizationStrategy strategy)
+    {
+        return ImmutableMap.builder()
+                .put("table", table)
+                .put("predicate", predicate.canonicalize(ignored -> false))
+                .build();
+    }
+}
diff --git a/presto-main-base/src/test/java/com/facebook/presto/connector/tvf/TestingTableFunctions.java b/presto-main-base/src/test/java/com/facebook/presto/connector/tvf/TestingTableFunctions.java
new file mode 100644
index 0000000000000..f2831f4bf10e1
--- /dev/null
+++ b/presto-main-base/src/test/java/com/facebook/presto/connector/tvf/TestingTableFunctions.java
@@ -0,0 +1,337 @@
+/*
+ * Licensed 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 com.facebook.presto.connector.tvf;
+
+import com.facebook.presto.common.predicate.TupleDomain;
+import com.facebook.presto.spi.ConnectorSession;
+import com.facebook.presto.spi.SchemaTableName;
+import com.facebook.presto.spi.connector.ConnectorTransactionHandle;
+import com.facebook.presto.spi.function.SchemaFunctionName;
+import com.facebook.presto.spi.function.table.AbstractConnectorTableFunction;
+import com.facebook.presto.spi.function.table.Argument;
+import com.facebook.presto.spi.function.table.ConnectorTableFunctionHandle;
+import com.facebook.presto.spi.function.table.Descriptor;
+import com.facebook.presto.spi.function.table.DescriptorArgumentSpecification;
+import com.facebook.presto.spi.function.table.ReturnTypeSpecification;
+import com.facebook.presto.spi.function.table.ScalarArgument;
+import com.facebook.presto.spi.function.table.ScalarArgumentSpecification;
+import com.facebook.presto.spi.function.table.TableArgumentSpecification;
+import com.facebook.presto.spi.function.table.TableFunctionAnalysis;
+import com.fasterxml.jackson.annotation.JsonCreator;
+import com.fasterxml.jackson.annotation.JsonInclude;
+import com.fasterxml.jackson.annotation.JsonProperty;
+import com.google.common.collect.ImmutableList;
+import io.airlift.slice.Slice;
+
+import java.util.Arrays;
+import java.util.Map;
+import java.util.Optional;
+
+import static com.facebook.presto.common.type.BigintType.BIGINT;
+import static com.facebook.presto.common.type.BooleanType.BOOLEAN;
+import static com.facebook.presto.common.type.IntegerType.INTEGER;
+import static com.facebook.presto.common.type.VarcharType.VARCHAR;
+import static com.facebook.presto.spi.function.table.ReturnTypeSpecification.GenericTable.GENERIC_TABLE;
+import static io.airlift.slice.Slices.utf8Slice;
+import static java.util.Objects.requireNonNull;
+
+public class TestingTableFunctions
+{
+    private static final String SCHEMA_NAME = "system";
+    private static final String TABLE_NAME = "table";
+    private static final String COLUMN_NAME = "column";
+    private static final ConnectorTableFunctionHandle HANDLE = new TestingTableFunctionPushdownHandle();
+    private static final TableFunctionAnalysis ANALYSIS = TableFunctionAnalysis.builder()
+            .handle(HANDLE)
+            .returnedType(new Descriptor(ImmutableList.of(new Descriptor.Field(COLUMN_NAME, Optional.of(BOOLEAN)))))
+            .build();
+
+    public static class TestConnectorTableFunction
+            extends AbstractConnectorTableFunction
+    {
+        private static final String TEST_FUNCTION = "test_function";
+
+        public TestConnectorTableFunction()
+        {
+            super(SCHEMA_NAME, TEST_FUNCTION, ImmutableList.of(), ReturnTypeSpecification.GenericTable.GENERIC_TABLE);
+        }
+
+        @Override
+        public TableFunctionAnalysis analyze(ConnectorSession session, ConnectorTransactionHandle transaction, Map arguments)
+        {
+            return TableFunctionAnalysis.builder()
+                    .handle(new TestingTableFunctionHandle(new SchemaFunctionName(SCHEMA_NAME, TEST_FUNCTION)))
+                    .returnedType(new Descriptor(ImmutableList.of(new Descriptor.Field("c1", Optional.of(BOOLEAN)))))
+                    .build();
+        }
+    }
+
+    public static class TestConnectorTableFunction2
+            extends AbstractConnectorTableFunction
+    {
+        private static final String TEST_FUNCTION_2 = "test_function2";
+
+        public TestConnectorTableFunction2()
+        {
+            super(SCHEMA_NAME, TEST_FUNCTION_2, ImmutableList.of(), ReturnTypeSpecification.OnlyPassThrough.ONLY_PASS_THROUGH);
+        }
+
+        @Override
+        public TableFunctionAnalysis analyze(ConnectorSession session, ConnectorTransactionHandle transaction, Map arguments)
+        {
+            return null;
+        }
+    }
+
+    public static class NullArgumentsTableFunction
+            extends AbstractConnectorTableFunction
+    {
+        private static final String NULL_ARGUMENTS_FUNCTION = "null_arguments_function";
+
+        public NullArgumentsTableFunction()
+        {
+            super(SCHEMA_NAME, NULL_ARGUMENTS_FUNCTION, null, ReturnTypeSpecification.OnlyPassThrough.ONLY_PASS_THROUGH);
+        }
+
+        @Override
+        public TableFunctionAnalysis analyze(ConnectorSession session, ConnectorTransactionHandle transaction, Map arguments)
+        {
+            return null;
+        }
+    }
+
+    public static class DuplicateArgumentsTableFunction
+            extends AbstractConnectorTableFunction
+    {
+        private static final String DUPLICATE_ARGUMENTS_FUNCTION = "duplicate_arguments_function";
+        public DuplicateArgumentsTableFunction()
+        {
+            super(
+                    SCHEMA_NAME,
+                    DUPLICATE_ARGUMENTS_FUNCTION,
+                    ImmutableList.of(
+                            ScalarArgumentSpecification.builder().name("a").type(INTEGER).build(),
+                            ScalarArgumentSpecification.builder().name("a").type(INTEGER).build()),
+                    ReturnTypeSpecification.OnlyPassThrough.ONLY_PASS_THROUGH);
+        }
+
+        @Override
+        public TableFunctionAnalysis analyze(ConnectorSession session, ConnectorTransactionHandle transaction, Map arguments)
+        {
+            return null;
+        }
+    }
+
+    public static class MultipleRSTableFunction
+            extends AbstractConnectorTableFunction
+    {
+        private static final String MULTIPLE_SOURCES_FUNCTION = "multiple_sources_function";
+        public MultipleRSTableFunction()
+        {
+            super(
+                    SCHEMA_NAME,
+                    MULTIPLE_SOURCES_FUNCTION,
+                    ImmutableList.of(TableArgumentSpecification.builder().name("t").rowSemantics().build(),
+                            TableArgumentSpecification.builder().name("t2").rowSemantics().build()),
+                    ReturnTypeSpecification.OnlyPassThrough.ONLY_PASS_THROUGH);
+        }
+
+        @Override
+        public TableFunctionAnalysis analyze(ConnectorSession session, ConnectorTransactionHandle transaction, Map arguments)
+        {
+            return null;
+        }
+    }
+
+    /**
+     * A table function returning a table with single empty column of type BOOLEAN.
+     * The argument `COLUMN` is the column name.
+     * The argument `IGNORED` is ignored.
+     * Both arguments are optional.
+     */
+    public static class SimpleTableFunction
+            extends AbstractConnectorTableFunction
+    {
+        private static final String FUNCTION_NAME = "simple_table_function";
+        private static final String TABLE_NAME = "simple_table";
+
+        public SimpleTableFunction()
+        {
+            super(
+                    SCHEMA_NAME,
+                    FUNCTION_NAME,
+                    Arrays.asList(
+                            ScalarArgumentSpecification.builder()
+                                    .name("COLUMN")
+                                    .type(VARCHAR)
+                                    .defaultValue(utf8Slice("col"))
+                                    .build(),
+                            ScalarArgumentSpecification.builder()
+                                    .name("IGNORED")
+                                    .type(BIGINT)
+                                    .defaultValue(0L)
+                                    .build()),
+                    GENERIC_TABLE);
+        }
+
+        @Override
+        public TableFunctionAnalysis analyze(ConnectorSession session, ConnectorTransactionHandle transaction, Map arguments)
+        {
+            ScalarArgument argument = (ScalarArgument) arguments.get("COLUMN");
+            String columnName = ((Slice) argument.getValue()).toStringUtf8();
+
+            return TableFunctionAnalysis.builder()
+                    .handle(new SimpleTableFunctionHandle(getSchema(), TABLE_NAME, columnName))
+                    .returnedType(new Descriptor(ImmutableList.of(new Descriptor.Field(columnName, Optional.of(BOOLEAN)))))
+                    .build();
+        }
+
+        public static class SimpleTableFunctionHandle
+                implements ConnectorTableFunctionHandle
+        {
+            private final TestTVFConnectorTableHandle tableHandle;
+
+            public SimpleTableFunctionHandle(String schema, String table, String column)
+            {
+                this.tableHandle = new TestTVFConnectorTableHandle(
+                        new SchemaTableName(schema, table),
+                        Optional.of(ImmutableList.of(new TestTVFConnectorColumnHandle(column, BOOLEAN))),
+                        TupleDomain.all());
+            }
+
+            public TestTVFConnectorTableHandle getTableHandle()
+            {
+                return tableHandle;
+            }
+        }
+    }
+
+    public static class TwoScalarArgumentsFunction
+            extends AbstractConnectorTableFunction
+    {
+        public TwoScalarArgumentsFunction()
+        {
+            super(
+                    SCHEMA_NAME,
+                    "two_arguments_function",
+                    ImmutableList.of(
+                            ScalarArgumentSpecification.builder()
+                                    .name("TEXT")
+                                    .type(VARCHAR)
+                                    .build(),
+                            ScalarArgumentSpecification.builder()
+                                    .name("NUMBER")
+                                    .type(BIGINT)
+                                    .defaultValue(null)
+                                    .build()),
+                    GENERIC_TABLE);
+        }
+
+        @Override
+        public TableFunctionAnalysis analyze(ConnectorSession session, ConnectorTransactionHandle transaction, Map arguments)
+        {
+            return ANALYSIS;
+        }
+    }
+
+    public static class TableArgumentFunction
+            extends AbstractConnectorTableFunction
+    {
+        public static final String FUNCTION_NAME = "table_argument_function";
+
+        public TableArgumentFunction()
+        {
+            super(
+                    SCHEMA_NAME,
+                    FUNCTION_NAME,
+                    ImmutableList.of(
+                            TableArgumentSpecification.builder()
+                                    .name("INPUT")
+                                    .keepWhenEmpty()
+                                    .build()),
+                    GENERIC_TABLE);
+        }
+
+        @Override
+        public TableFunctionAnalysis analyze(ConnectorSession session, ConnectorTransactionHandle transaction, Map arguments)
+        {
+            return TableFunctionAnalysis.builder()
+                    .handle(new TestingTableFunctionHandle(new SchemaFunctionName(SCHEMA_NAME, FUNCTION_NAME)))
+                    .returnedType(new Descriptor(ImmutableList.of(new Descriptor.Field(COLUMN_NAME, Optional.of(BOOLEAN)))))
+                    .requiredColumns("INPUT", ImmutableList.of(0))
+                    .build();
+        }
+    }
+
+    public static class DescriptorArgumentFunction
+            extends AbstractConnectorTableFunction
+    {
+        public DescriptorArgumentFunction()
+        {
+            super(
+                    SCHEMA_NAME,
+                    "descriptor_argument_function",
+                    ImmutableList.of(
+                            DescriptorArgumentSpecification.builder()
+                                    .name("SCHEMA")
+                                    .defaultValue(null)
+                                    .build()),
+                    GENERIC_TABLE);
+        }
+
+        @Override
+        public TableFunctionAnalysis analyze(ConnectorSession session, ConnectorTransactionHandle transaction, Map arguments)
+        {
+            return ANALYSIS;
+        }
+    }
+
+    public static class TestingTableFunctionPushdownHandle
+            implements ConnectorTableFunctionHandle
+    {
+        private final TestTVFConnectorTableHandle tableHandle;
+
+        public TestingTableFunctionPushdownHandle()
+        {
+            this.tableHandle = new TestTVFConnectorTableHandle(
+                    new SchemaTableName(SCHEMA_NAME, TABLE_NAME),
+                    Optional.of(ImmutableList.of(new TestTVFConnectorColumnHandle(COLUMN_NAME, BOOLEAN))),
+                    TupleDomain.all());
+        }
+
+        public TestTVFConnectorTableHandle getTableHandle()
+        {
+            return tableHandle;
+        }
+    }
+
+    @JsonInclude(JsonInclude.Include.ALWAYS)
+    public static class TestingTableFunctionHandle
+            implements ConnectorTableFunctionHandle
+    {
+        private final SchemaFunctionName schemaFunctionName;
+
+        @JsonCreator
+        public TestingTableFunctionHandle(@JsonProperty("schemaFunctionName") SchemaFunctionName schemaFunctionName)
+        {
+            this.schemaFunctionName = requireNonNull(schemaFunctionName, "schemaFunctionName is null");
+        }
+
+        @JsonProperty
+        public SchemaFunctionName getSchemaFunctionName()
+        {
+            return schemaFunctionName;
+        }
+    }
+}
diff --git a/presto-main-base/src/test/java/com/facebook/presto/execution/TestClusterOverloadConfig.java b/presto-main-base/src/test/java/com/facebook/presto/execution/TestClusterOverloadConfig.java
new file mode 100644
index 0000000000000..b28aee410fb20
--- /dev/null
+++ b/presto-main-base/src/test/java/com/facebook/presto/execution/TestClusterOverloadConfig.java
@@ -0,0 +1,57 @@
+/*
+ * Licensed 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 com.facebook.presto.execution;
+
+import com.facebook.airlift.configuration.testing.ConfigAssertions;
+import com.google.common.collect.ImmutableMap;
+import org.testng.annotations.Test;
+
+import java.util.Map;
+
+import static com.facebook.presto.execution.ClusterOverloadConfig.OVERLOAD_POLICY_CNT_BASED;
+
+public class TestClusterOverloadConfig
+{
+    @Test
+    public void testDefaults()
+    {
+        ConfigAssertions.assertRecordedDefaults(ConfigAssertions.recordDefaults(ClusterOverloadConfig.class)
+                .setClusterOverloadThrottlingEnabled(false)
+                .setAllowedOverloadWorkersPct(0.01)
+                .setAllowedOverloadWorkersCnt(0)
+                .setOverloadPolicyType(OVERLOAD_POLICY_CNT_BASED)
+                .setOverloadCheckCacheTtlInSecs(5));
+    }
+
+    @Test
+    public void testExplicitPropertyMappings()
+    {
+        Map properties = new ImmutableMap.Builder()
+                .put("cluster-overload.enable-throttling", "true")
+                .put("cluster-overload.allowed-overload-workers-pct", "0.05")
+                .put("cluster-overload.allowed-overload-workers-cnt", "5")
+                .put("cluster-overload.overload-policy-type", "overload_worker_pct_based_throttling")
+                .put("cluster.overload-check-cache-ttl-secs", "10")
+                .build();
+
+        ClusterOverloadConfig expected = new ClusterOverloadConfig()
+                .setClusterOverloadThrottlingEnabled(true)
+                .setAllowedOverloadWorkersPct(0.05)
+                .setAllowedOverloadWorkersCnt(5)
+                .setOverloadPolicyType("overload_worker_pct_based_throttling")
+                .setOverloadCheckCacheTtlInSecs(10);
+
+        ConfigAssertions.assertFullMapping(properties, expected);
+    }
+}
diff --git a/presto-main-base/src/test/java/com/facebook/presto/execution/TestOutput.java b/presto-main-base/src/test/java/com/facebook/presto/execution/TestOutput.java
index a427cf4f4878b..8e715deea4376 100644
--- a/presto-main-base/src/test/java/com/facebook/presto/execution/TestOutput.java
+++ b/presto-main-base/src/test/java/com/facebook/presto/execution/TestOutput.java
@@ -14,8 +14,12 @@
 package com.facebook.presto.execution;
 
 import com.facebook.airlift.json.JsonCodec;
+import com.facebook.presto.common.QualifiedObjectName;
+import com.facebook.presto.common.SourceColumn;
 import com.facebook.presto.spi.ConnectorId;
+import com.facebook.presto.spi.eventlistener.OutputColumnMetadata;
 import com.google.common.collect.ImmutableList;
+import com.google.common.collect.ImmutableSet;
 import org.testng.annotations.Test;
 
 import java.util.Optional;
@@ -37,7 +41,10 @@ public void testRoundTrip()
                 EMPTY_COMMIT_OUTPUT,
                 Optional.of(
                         ImmutableList.of(
-                                new Column("column", "type"))));
+                                new OutputColumnMetadata(
+                                        "column", "type",
+                                        ImmutableSet.of(
+                                                new SourceColumn(QualifiedObjectName.valueOf("catalog.schema.table"), "column"))))));
 
         String json = codec.toJson(expected);
         Output actual = codec.fromJson(json);
diff --git a/presto-main-base/src/test/java/com/facebook/presto/execution/resourceGroups/BenchmarkResourceGroup.java b/presto-main-base/src/test/java/com/facebook/presto/execution/resourceGroups/BenchmarkResourceGroup.java
index dfa9235e2f4f7..1ab8ba0bc30e7 100644
--- a/presto-main-base/src/test/java/com/facebook/presto/execution/resourceGroups/BenchmarkResourceGroup.java
+++ b/presto-main-base/src/test/java/com/facebook/presto/execution/resourceGroups/BenchmarkResourceGroup.java
@@ -14,9 +14,13 @@
 package com.facebook.presto.execution.resourceGroups;
 
 import com.facebook.airlift.units.DataSize;
+import com.facebook.presto.execution.ClusterOverloadConfig;
 import com.facebook.presto.execution.MockManagedQueryExecution;
 import com.facebook.presto.execution.resourceGroups.InternalResourceGroup.RootInternalResourceGroup;
+import com.facebook.presto.execution.scheduler.clusterOverload.ClusterOverloadPolicy;
+import com.facebook.presto.execution.scheduler.clusterOverload.ClusterResourceChecker;
 import com.facebook.presto.metadata.InMemoryNodeManager;
+import com.facebook.presto.metadata.InternalNodeManager;
 import org.openjdk.jmh.annotations.Benchmark;
 import org.openjdk.jmh.annotations.BenchmarkMode;
 import org.openjdk.jmh.annotations.Fork;
@@ -73,7 +77,7 @@ public static class BenchmarkData
         @Setup
         public void setup()
         {
-            root = new RootInternalResourceGroup("root", (group, export) -> {}, executor, ignored -> Optional.empty(), rg -> false, new InMemoryNodeManager());
+            root = new RootInternalResourceGroup("root", (group, export) -> {}, executor, ignored -> Optional.empty(), rg -> false, new InMemoryNodeManager(), createClusterResourceChecker());
             root.setSoftMemoryLimit(new DataSize(1, MEGABYTE));
             root.setMaxQueuedQueries(queries);
             root.setHardConcurrencyLimit(queries);
@@ -89,6 +93,31 @@ public void setup()
             }
         }
 
+        private ClusterResourceChecker createClusterResourceChecker()
+        {
+            // Create a mock cluster overload policy that never reports overload
+            ClusterOverloadPolicy mockPolicy = new ClusterOverloadPolicy()
+            {
+                @Override
+                public boolean isClusterOverloaded(InternalNodeManager nodeManager)
+                {
+                    return false; // Never overloaded for benchmarks
+                }
+
+                @Override
+                public String getName()
+                {
+                    return "benchmark-policy";
+                }
+            };
+
+            // Create a config with throttling disabled for benchmarks
+            ClusterOverloadConfig config = new ClusterOverloadConfig()
+                    .setClusterOverloadThrottlingEnabled(false);
+
+            return new ClusterResourceChecker(mockPolicy, config, new InMemoryNodeManager());
+        }
+
         @TearDown
         public void tearDown()
         {
diff --git a/presto-main-base/src/test/java/com/facebook/presto/execution/resourceGroups/TestInternalResourceGroupManager.java b/presto-main-base/src/test/java/com/facebook/presto/execution/resourceGroups/TestInternalResourceGroupManager.java
index bc4cf99103531..257d957f8e578 100644
--- a/presto-main-base/src/test/java/com/facebook/presto/execution/resourceGroups/TestInternalResourceGroupManager.java
+++ b/presto-main-base/src/test/java/com/facebook/presto/execution/resourceGroups/TestInternalResourceGroupManager.java
@@ -15,8 +15,11 @@
 package com.facebook.presto.execution.resourceGroups;
 
 import com.facebook.airlift.node.NodeInfo;
+import com.facebook.presto.execution.ClusterOverloadConfig;
 import com.facebook.presto.execution.MockManagedQueryExecution;
 import com.facebook.presto.execution.QueryManagerConfig;
+import com.facebook.presto.execution.scheduler.clusterOverload.ClusterResourceChecker;
+import com.facebook.presto.execution.scheduler.clusterOverload.CpuMemoryOverloadPolicy;
 import com.facebook.presto.metadata.InMemoryNodeManager;
 import com.facebook.presto.server.ServerConfig;
 import com.facebook.presto.spi.PrestoException;
@@ -32,8 +35,7 @@ public class TestInternalResourceGroupManager
     @Test(expectedExceptions = PrestoException.class, expectedExceptionsMessageRegExp = ".*Presto server is still initializing.*")
     public void testQueryFailsWithInitializingConfigurationManager()
     {
-        InternalResourceGroupManager> internalResourceGroupManager = new InternalResourceGroupManager<>((poolId, listener) -> {},
-                new QueryManagerConfig(), new NodeInfo("test"), new MBeanExporter(new TestingMBeanServer()), () -> null, new ServerConfig(), new InMemoryNodeManager());
+        InternalResourceGroupManager> internalResourceGroupManager = new InternalResourceGroupManager<>((poolId, listener) -> {}, new QueryManagerConfig(), new NodeInfo("test"), new MBeanExporter(new TestingMBeanServer()), () -> null, new ServerConfig(), new InMemoryNodeManager(), new ClusterResourceChecker(new CpuMemoryOverloadPolicy(new ClusterOverloadConfig()), new ClusterOverloadConfig(), new InMemoryNodeManager()));
         internalResourceGroupManager.submit(new MockManagedQueryExecution(0), new SelectionContext<>(new ResourceGroupId("global"), ImmutableMap.of()), command -> {});
     }
 
@@ -42,7 +44,7 @@ public void testQuerySucceedsWhenConfigurationManagerLoaded()
             throws Exception
     {
         InternalResourceGroupManager> internalResourceGroupManager = new InternalResourceGroupManager<>((poolId, listener) -> {},
-                new QueryManagerConfig(), new NodeInfo("test"), new MBeanExporter(new TestingMBeanServer()), () -> null, new ServerConfig(), new InMemoryNodeManager());
+                new QueryManagerConfig(), new NodeInfo("test"), new MBeanExporter(new TestingMBeanServer()), () -> null, new ServerConfig(), new InMemoryNodeManager(), new ClusterResourceChecker(new CpuMemoryOverloadPolicy(new ClusterOverloadConfig()), new ClusterOverloadConfig(), new InMemoryNodeManager()));
         internalResourceGroupManager.loadConfigurationManager();
         internalResourceGroupManager.submit(new MockManagedQueryExecution(0), new SelectionContext<>(new ResourceGroupId("global"), ImmutableMap.of()), command -> {});
     }
diff --git a/presto-main-base/src/test/java/com/facebook/presto/execution/resourceGroups/TestResourceGroups.java b/presto-main-base/src/test/java/com/facebook/presto/execution/resourceGroups/TestResourceGroups.java
index 2b0eb9ac7d299..5890a210c44e1 100644
--- a/presto-main-base/src/test/java/com/facebook/presto/execution/resourceGroups/TestResourceGroups.java
+++ b/presto-main-base/src/test/java/com/facebook/presto/execution/resourceGroups/TestResourceGroups.java
@@ -15,8 +15,11 @@
 
 import com.facebook.airlift.units.DataSize;
 import com.facebook.airlift.units.Duration;
+import com.facebook.presto.execution.ClusterOverloadConfig;
 import com.facebook.presto.execution.MockManagedQueryExecution;
 import com.facebook.presto.execution.resourceGroups.InternalResourceGroup.RootInternalResourceGroup;
+import com.facebook.presto.execution.scheduler.clusterOverload.ClusterOverloadPolicy;
+import com.facebook.presto.execution.scheduler.clusterOverload.ClusterResourceChecker;
 import com.facebook.presto.metadata.InMemoryNodeManager;
 import com.facebook.presto.metadata.InternalNode;
 import com.facebook.presto.metadata.InternalNodeManager;
@@ -69,7 +72,7 @@ public class TestResourceGroups
     @Test(timeOut = 10_000)
     public void testQueueFull()
     {
-        RootInternalResourceGroup root = new RootInternalResourceGroup("root", (group, export) -> {}, directExecutor(), ignored -> Optional.empty(), rg -> false, createNodeManager());
+        RootInternalResourceGroup root = new RootInternalResourceGroup("root", (group, export) -> {}, directExecutor(), ignored -> Optional.empty(), rg -> false, createNodeManager(), createClusterResourceChecker());
         root.setSoftMemoryLimit(new DataSize(1, MEGABYTE));
         root.setMaxQueuedQueries(1);
         root.setHardConcurrencyLimit(1);
@@ -91,7 +94,7 @@ public void testQueueFull()
     @Test(timeOut = 10_000)
     public void testFairEligibility()
     {
-        RootInternalResourceGroup root = new RootInternalResourceGroup("root", (group, export) -> {}, directExecutor(), ignored -> Optional.empty(), rg -> false, createNodeManager());
+        RootInternalResourceGroup root = new RootInternalResourceGroup("root", (group, export) -> {}, directExecutor(), ignored -> Optional.empty(), rg -> false, createNodeManager(), createClusterResourceChecker());
         root.setSoftMemoryLimit(new DataSize(1, MEGABYTE));
         root.setMaxQueuedQueries(4);
         root.setHardConcurrencyLimit(1);
@@ -151,7 +154,7 @@ public void testFairEligibility()
     @Test
     public void testSetSchedulingPolicy()
     {
-        RootInternalResourceGroup root = new RootInternalResourceGroup("root", (group, export) -> {}, directExecutor(), ignored -> Optional.empty(), rg -> false, createNodeManager());
+        RootInternalResourceGroup root = new RootInternalResourceGroup("root", (group, export) -> {}, directExecutor(), ignored -> Optional.empty(), rg -> false, createNodeManager(), createClusterResourceChecker());
         root.setSoftMemoryLimit(new DataSize(1, MEGABYTE));
         root.setMaxQueuedQueries(4);
         root.setHardConcurrencyLimit(1);
@@ -197,7 +200,7 @@ public void testSetSchedulingPolicy()
     @Test(timeOut = 10_000)
     public void testFairQueuing()
     {
-        RootInternalResourceGroup root = new RootInternalResourceGroup("root", (group, export) -> {}, directExecutor(), ignored -> Optional.empty(), rg -> false, createNodeManager());
+        RootInternalResourceGroup root = new RootInternalResourceGroup("root", (group, export) -> {}, directExecutor(), ignored -> Optional.empty(), rg -> false, createNodeManager(), createClusterResourceChecker());
         root.setSoftMemoryLimit(new DataSize(1, MEGABYTE));
         root.setMaxQueuedQueries(4);
         root.setHardConcurrencyLimit(1);
@@ -243,7 +246,7 @@ public void testFairQueuing()
     @Test(timeOut = 10_000)
     public void testMemoryLimit()
     {
-        RootInternalResourceGroup root = new RootInternalResourceGroup("root", (group, export) -> {}, directExecutor(), ignored -> Optional.empty(), rg -> false, createNodeManager());
+        RootInternalResourceGroup root = new RootInternalResourceGroup("root", (group, export) -> {}, directExecutor(), ignored -> Optional.empty(), rg -> false, createNodeManager(), createClusterResourceChecker());
         root.setSoftMemoryLimit(new DataSize(1, BYTE));
         root.setMaxQueuedQueries(4);
         root.setHardConcurrencyLimit(3);
@@ -271,7 +274,7 @@ public void testMemoryLimit()
     @Test
     public void testSubgroupMemoryLimit()
     {
-        RootInternalResourceGroup root = new RootInternalResourceGroup("root", (group, export) -> {}, directExecutor(), ignored -> Optional.empty(), rg -> false, createNodeManager());
+        RootInternalResourceGroup root = new RootInternalResourceGroup("root", (group, export) -> {}, directExecutor(), ignored -> Optional.empty(), rg -> false, createNodeManager(), createClusterResourceChecker());
         root.setSoftMemoryLimit(new DataSize(10, BYTE));
         root.setMaxQueuedQueries(4);
         root.setHardConcurrencyLimit(3);
@@ -304,7 +307,7 @@ public void testSubgroupMemoryLimit()
     @Test(timeOut = 10_000)
     public void testSoftCpuLimit()
     {
-        RootInternalResourceGroup root = new RootInternalResourceGroup("root", (group, export) -> {}, directExecutor(), ignored -> Optional.empty(), rg -> false, createNodeManager());
+        RootInternalResourceGroup root = new RootInternalResourceGroup("root", (group, export) -> {}, directExecutor(), ignored -> Optional.empty(), rg -> false, createNodeManager(), createClusterResourceChecker());
         root.setSoftMemoryLimit(new DataSize(1, BYTE));
         root.setSoftCpuLimit(new Duration(1, SECONDS));
         root.setHardCpuLimit(new Duration(2, SECONDS));
@@ -341,7 +344,7 @@ public void testSoftCpuLimit()
     @Test(timeOut = 10_000)
     public void testPerWorkerQueryLimit()
     {
-        RootInternalResourceGroup root = new RootInternalResourceGroup("root", (group, export) -> {}, directExecutor(), ignored -> Optional.empty(), rg -> false, createNodeManager());
+        RootInternalResourceGroup root = new RootInternalResourceGroup("root", (group, export) -> {}, directExecutor(), ignored -> Optional.empty(), rg -> false, createNodeManager(), createClusterResourceChecker());
         root.setWorkersPerQueryLimit(5);
         root.setMaxQueuedQueries(2);
         root.setHardConcurrencyLimit(2);
@@ -374,7 +377,7 @@ public void testPerWorkerQueryLimit()
     @Test(timeOut = 10_000)
     public void testPerWorkerQueryLimitMultipleGroups()
     {
-        RootInternalResourceGroup root = new RootInternalResourceGroup("root", (group, export) -> {}, directExecutor(), ignored -> Optional.empty(), rg -> false, createNodeManager());
+        RootInternalResourceGroup root = new RootInternalResourceGroup("root", (group, export) -> {}, directExecutor(), ignored -> Optional.empty(), rg -> false, createNodeManager(), createClusterResourceChecker());
         root.setWorkersPerQueryLimit(5);
         root.setMaxQueuedQueries(5);
         root.setHardConcurrencyLimit(2);
@@ -417,7 +420,7 @@ public void testPerWorkerQueryLimitMultipleGroups()
     @Test(timeOut = 10_000)
     public void testHardCpuLimit()
     {
-        RootInternalResourceGroup root = new RootInternalResourceGroup("root", (group, export) -> {}, directExecutor(), ignored -> Optional.empty(), rg -> false, createNodeManager());
+        RootInternalResourceGroup root = new RootInternalResourceGroup("root", (group, export) -> {}, directExecutor(), ignored -> Optional.empty(), rg -> false, createNodeManager(), createClusterResourceChecker());
         root.setSoftMemoryLimit(new DataSize(1, BYTE));
         root.setHardCpuLimit(new Duration(1, SECONDS));
         root.setCpuQuotaGenerationMillisPerSecond(2000);
@@ -444,7 +447,7 @@ public void testHardCpuLimit()
     @Test(timeOut = 10_000)
     public void testPriorityScheduling()
     {
-        RootInternalResourceGroup root = new RootInternalResourceGroup("root", (group, export) -> {}, directExecutor(), ignored -> Optional.empty(), rg -> false, createNodeManager());
+        RootInternalResourceGroup root = new RootInternalResourceGroup("root", (group, export) -> {}, directExecutor(), ignored -> Optional.empty(), rg -> false, createNodeManager(), createClusterResourceChecker());
         root.setSoftMemoryLimit(new DataSize(1, MEGABYTE));
         root.setMaxQueuedQueries(100);
         // Start with zero capacity, so that nothing starts running until we've added all the queries
@@ -494,7 +497,7 @@ public void testPriorityScheduling()
     @Test(timeOut = 20_000)
     public void testWeightedScheduling()
     {
-        RootInternalResourceGroup root = new RootInternalResourceGroup("root", (group, export) -> {}, directExecutor(), ignored -> Optional.empty(), rg -> false, createNodeManager());
+        RootInternalResourceGroup root = new RootInternalResourceGroup("root", (group, export) -> {}, directExecutor(), ignored -> Optional.empty(), rg -> false, createNodeManager(), createClusterResourceChecker());
         root.setSoftMemoryLimit(new DataSize(1, MEGABYTE));
         root.setMaxQueuedQueries(4);
         // Start with zero capacity, so that nothing starts running until we've added all the queries
@@ -543,7 +546,7 @@ public void testWeightedScheduling()
     @Test(timeOut = 30_000)
     public void testWeightedFairScheduling()
     {
-        RootInternalResourceGroup root = new RootInternalResourceGroup("root", (group, export) -> {}, directExecutor(), ignored -> Optional.empty(), rg -> false, createNodeManager());
+        RootInternalResourceGroup root = new RootInternalResourceGroup("root", (group, export) -> {}, directExecutor(), ignored -> Optional.empty(), rg -> false, createNodeManager(), createClusterResourceChecker());
         root.setSoftMemoryLimit(new DataSize(1, MEGABYTE));
         root.setMaxQueuedQueries(50);
         // Start with zero capacity, so that nothing starts running until we've added all the queries
@@ -586,7 +589,7 @@ public void testWeightedFairScheduling()
     @Test(timeOut = 10_000)
     public void testWeightedFairSchedulingEqualWeights()
     {
-        RootInternalResourceGroup root = new RootInternalResourceGroup("root", (group, export) -> {}, directExecutor(), ignored -> Optional.empty(), rg -> false, createNodeManager());
+        RootInternalResourceGroup root = new RootInternalResourceGroup("root", (group, export) -> {}, directExecutor(), ignored -> Optional.empty(), rg -> false, createNodeManager(), createClusterResourceChecker());
         root.setSoftMemoryLimit(new DataSize(1, MEGABYTE));
         root.setMaxQueuedQueries(50);
         // Start with zero capacity, so that nothing starts running until we've added all the queries
@@ -645,7 +648,7 @@ public void testWeightedFairSchedulingEqualWeights()
     @Test(timeOut = 20_000)
     public void testWeightedFairSchedulingNoStarvation()
     {
-        RootInternalResourceGroup root = new RootInternalResourceGroup("root", (group, export) -> {}, directExecutor(), ignored -> Optional.empty(), rg -> false, createNodeManager());
+        RootInternalResourceGroup root = new RootInternalResourceGroup("root", (group, export) -> {}, directExecutor(), ignored -> Optional.empty(), rg -> false, createNodeManager(), createClusterResourceChecker());
         root.setSoftMemoryLimit(new DataSize(1, MEGABYTE));
         root.setMaxQueuedQueries(50);
         // Start with zero capacity, so that nothing starts running until we've added all the queries
@@ -686,7 +689,7 @@ public void testWeightedFairSchedulingNoStarvation()
     @Test
     public void testGetInfo()
     {
-        RootInternalResourceGroup root = new RootInternalResourceGroup("root", (group, export) -> {}, directExecutor(), ignored -> Optional.empty(), rg -> false, createNodeManager());
+        RootInternalResourceGroup root = new RootInternalResourceGroup("root", (group, export) -> {}, directExecutor(), ignored -> Optional.empty(), rg -> false, createNodeManager(), createClusterResourceChecker());
         root.setSoftMemoryLimit(new DataSize(1, MEGABYTE));
         root.setMaxQueuedQueries(40);
         // Start with zero capacity, so that nothing starts running until we've added all the queries
@@ -776,7 +779,7 @@ public void testGetInfo()
     @Test
     public void testGetResourceGroupStateInfo()
     {
-        RootInternalResourceGroup root = new RootInternalResourceGroup("root", (group, export) -> {}, directExecutor(), ignored -> Optional.empty(), rg -> false, createNodeManager());
+        RootInternalResourceGroup root = new RootInternalResourceGroup("root", (group, export) -> {}, directExecutor(), ignored -> Optional.empty(), rg -> false, createNodeManager(), createClusterResourceChecker());
         root.setSoftMemoryLimit(new DataSize(1, GIGABYTE));
         root.setMaxQueuedQueries(40);
         root.setHardConcurrencyLimit(10);
@@ -844,7 +847,7 @@ public void testGetResourceGroupStateInfo()
     @Test
     public void testGetStaticResourceGroupInfo()
     {
-        RootInternalResourceGroup root = new RootInternalResourceGroup("root", (group, export) -> {}, directExecutor(), ignored -> Optional.empty(), rg -> false, createNodeManager());
+        RootInternalResourceGroup root = new RootInternalResourceGroup("root", (group, export) -> {}, directExecutor(), ignored -> Optional.empty(), rg -> false, createNodeManager(), createClusterResourceChecker());
         root.setSoftMemoryLimit(new DataSize(1, GIGABYTE));
         root.setMaxQueuedQueries(100);
         root.setHardConcurrencyLimit(10);
@@ -921,7 +924,7 @@ private Optional getResourceGroupInfoForId(InternalResourceGr
     @Test
     public void testGetBlockedQueuedQueries()
     {
-        RootInternalResourceGroup root = new RootInternalResourceGroup("root", (group, export) -> {}, directExecutor(), ignored -> Optional.empty(), rg -> false, createNodeManager());
+        RootInternalResourceGroup root = new RootInternalResourceGroup("root", (group, export) -> {}, directExecutor(), ignored -> Optional.empty(), rg -> false, createNodeManager(), createClusterResourceChecker());
         root.setSoftMemoryLimit(new DataSize(1, MEGABYTE));
         root.setMaxQueuedQueries(40);
         // Start with zero capacity, so that nothing starts running until we've added all the queries
@@ -1072,4 +1075,27 @@ private InternalNodeManager createNodeManager()
                         false));
         return internalNodeManager;
     }
+
+    private ClusterResourceChecker createClusterResourceChecker()
+    {
+        ClusterOverloadPolicy mockPolicy = new ClusterOverloadPolicy()
+        {
+            @Override
+            public boolean isClusterOverloaded(InternalNodeManager nodeManager)
+            {
+                return false;
+            }
+
+            @Override
+            public String getName()
+            {
+                return "test-policy";
+            }
+        };
+
+        ClusterOverloadConfig config = new ClusterOverloadConfig()
+                .setClusterOverloadThrottlingEnabled(false);
+
+        return new ClusterResourceChecker(mockPolicy, config, createNodeManager());
+    }
 }
diff --git a/presto-main-base/src/test/java/com/facebook/presto/execution/scheduler/clusterOverload/TestClusterResourceChecker.java b/presto-main-base/src/test/java/com/facebook/presto/execution/scheduler/clusterOverload/TestClusterResourceChecker.java
new file mode 100644
index 0000000000000..1a5d9c4af8371
--- /dev/null
+++ b/presto-main-base/src/test/java/com/facebook/presto/execution/scheduler/clusterOverload/TestClusterResourceChecker.java
@@ -0,0 +1,305 @@
+/*
+ * Licensed 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 com.facebook.presto.execution.scheduler.clusterOverload;
+
+import com.facebook.presto.execution.ClusterOverloadConfig;
+import com.facebook.presto.metadata.AllNodes;
+import com.facebook.presto.metadata.InternalNode;
+import com.facebook.presto.metadata.InternalNodeManager;
+import com.facebook.presto.spi.ConnectorId;
+import com.facebook.presto.spi.NodeLoadMetrics;
+import com.facebook.presto.spi.NodeState;
+import org.testng.annotations.BeforeMethod;
+import org.testng.annotations.Test;
+
+import java.util.Collections;
+import java.util.Optional;
+import java.util.Set;
+import java.util.concurrent.atomic.AtomicInteger;
+import java.util.function.Consumer;
+
+import static org.testng.Assert.assertEquals;
+import static org.testng.Assert.assertFalse;
+import static org.testng.Assert.assertTrue;
+
+@Test(singleThreaded = true)
+public class TestClusterResourceChecker
+{
+    private static final int CACHE_TTL_SECS = 1;
+    private static final int SLEEP_BUFFER_MILLIS = 200;
+
+    private ClusterOverloadConfig config;
+    private TestingClusterOverloadPolicy clusterOverloadPolicy;
+    private ClusterResourceChecker clusterResourceChecker;
+    private TestingInternalNodeManager nodeManager;
+
+    @BeforeMethod
+    public void setUp()
+    {
+        config = new ClusterOverloadConfig()
+                .setOverloadCheckCacheTtlInSecs(CACHE_TTL_SECS)
+                .setClusterOverloadThrottlingEnabled(true);
+
+        clusterOverloadPolicy = new TestingClusterOverloadPolicy();
+        nodeManager = new TestingInternalNodeManager();
+
+        clusterResourceChecker = new ClusterResourceChecker(clusterOverloadPolicy, config, nodeManager);
+    }
+
+    public void testInitialState()
+    {
+        assertFalse(clusterResourceChecker.isClusterOverloaded());
+        assertEquals(clusterResourceChecker.getOverloadDetectionCount().getTotalCount(), 0);
+        assertEquals(clusterResourceChecker.getOverloadDurationMillis(), 0);
+        assertTrue(clusterResourceChecker.isClusterOverloadThrottlingEnabled());
+    }
+
+    @Test
+    public void testIsClusterCurrentlyOverloaded()
+    {
+        // Start the periodic task
+        clusterResourceChecker.start();
+
+        // Initially not overloaded
+        clusterOverloadPolicy.setOverloaded(false);
+        assertFalse(clusterResourceChecker.isClusterCurrentlyOverloaded());
+
+        // Wait for periodic check to update state
+        clusterOverloadPolicy.setOverloaded(true);
+        sleep((CACHE_TTL_SECS * 1000) + SLEEP_BUFFER_MILLIS);
+        assertTrue(clusterResourceChecker.isClusterCurrentlyOverloaded());
+        assertTrue(clusterResourceChecker.isClusterOverloaded());
+        assertEquals(clusterResourceChecker.getOverloadDetectionCount().getTotalCount(), 1);
+
+        // Set back to not overloaded
+        clusterOverloadPolicy.setOverloaded(false);
+        sleep((CACHE_TTL_SECS * 1000) + SLEEP_BUFFER_MILLIS);
+        assertFalse(clusterResourceChecker.isClusterCurrentlyOverloaded());
+        assertFalse(clusterResourceChecker.isClusterOverloaded());
+
+        // Stop the periodic task
+        clusterResourceChecker.stop();
+    }
+
+    @Test
+    public void testOverloadDurationMetric()
+    {
+        // Start the periodic task
+        clusterResourceChecker.start();
+
+        // Set to overloaded and wait for periodic check
+        clusterOverloadPolicy.setOverloaded(true);
+        sleep((CACHE_TTL_SECS * 1000) + SLEEP_BUFFER_MILLIS);
+        assertTrue(clusterResourceChecker.isClusterCurrentlyOverloaded());
+        assertTrue(clusterResourceChecker.isClusterOverloaded());
+
+        // Duration should be greater than 0
+        sleep(100);
+        assertTrue(clusterResourceChecker.getOverloadDurationMillis() > 0);
+
+        // Set back to not overloaded
+        clusterOverloadPolicy.setOverloaded(false);
+        sleep((CACHE_TTL_SECS * 1000) + SLEEP_BUFFER_MILLIS);
+        assertFalse(clusterResourceChecker.isClusterCurrentlyOverloaded());
+
+        // Duration should be 0 again
+        assertEquals(clusterResourceChecker.getOverloadDurationMillis(), 0);
+
+        // Stop the periodic task
+        clusterResourceChecker.stop();
+    }
+
+    @Test
+    public void testMultipleOverloadTransitions()
+    {
+        // Start the periodic task
+        clusterResourceChecker.start();
+
+        // First transition to overloaded
+        clusterOverloadPolicy.setOverloaded(true);
+        sleep((CACHE_TTL_SECS * 1000) + SLEEP_BUFFER_MILLIS);
+        assertTrue(clusterResourceChecker.isClusterCurrentlyOverloaded());
+        assertEquals(clusterResourceChecker.getOverloadDetectionCount().getTotalCount(), 1);
+
+        // Wait and transition back to not overloaded
+        clusterOverloadPolicy.setOverloaded(false);
+        sleep((CACHE_TTL_SECS * 1000) + SLEEP_BUFFER_MILLIS);
+        assertFalse(clusterResourceChecker.isClusterCurrentlyOverloaded());
+
+        // Second transition to overloaded
+        clusterOverloadPolicy.setOverloaded(true);
+        sleep((CACHE_TTL_SECS * 1000) + SLEEP_BUFFER_MILLIS);
+        assertTrue(clusterResourceChecker.isClusterCurrentlyOverloaded());
+        assertEquals(clusterResourceChecker.getOverloadDetectionCount().getTotalCount(), 2);
+
+        // Stop the periodic task
+        clusterResourceChecker.stop();
+    }
+
+    @Test
+    public void testClusterOverloadThrottlingEnabled()
+    {
+        // Default is enabled (set in setUp)
+        assertTrue(clusterResourceChecker.isClusterOverloadThrottlingEnabled());
+
+        // Create a new config with throttling disabled
+        ClusterOverloadConfig disabledConfig = new ClusterOverloadConfig()
+                .setOverloadCheckCacheTtlInSecs(CACHE_TTL_SECS)
+                .setClusterOverloadThrottlingEnabled(false);
+
+        // Create a new checker with throttling disabled
+        ClusterResourceChecker disabledChecker = new ClusterResourceChecker(clusterOverloadPolicy, disabledConfig, nodeManager);
+        assertFalse(disabledChecker.isClusterOverloadThrottlingEnabled());
+
+        // Even when cluster is overloaded, isClusterCurrentlyOverloaded should return false if throttling is disabled
+        clusterOverloadPolicy.setOverloaded(true);
+        assertFalse(disabledChecker.isClusterCurrentlyOverloaded());
+    }
+
+    private void sleep(long millis)
+    {
+        try {
+            Thread.sleep(millis);
+        }
+        catch (InterruptedException e) {
+            Thread.currentThread().interrupt();
+            throw new RuntimeException(e);
+        }
+    }
+
+    private static class TestingClusterOverloadPolicy
+            implements ClusterOverloadPolicy
+    {
+        private boolean overloaded;
+        private final AtomicInteger checkCount = new AtomicInteger();
+
+        @Override
+        public boolean isClusterOverloaded(InternalNodeManager nodeManager)
+        {
+            checkCount.incrementAndGet();
+            return overloaded;
+        }
+
+        @Override
+        public String getName()
+        {
+            return "test-policy";
+        }
+
+        public void setOverloaded(boolean overloaded)
+        {
+            this.overloaded = overloaded;
+        }
+
+        public int getCheckCount()
+        {
+            return checkCount.get();
+        }
+    }
+
+    private static class TestingInternalNodeManager
+            implements InternalNodeManager
+    {
+        @Override
+        public Set getNodes(NodeState state)
+        {
+            return Collections.emptySet();
+        }
+
+        @Override
+        public Set getActiveConnectorNodes(ConnectorId connectorId)
+        {
+            return Collections.emptySet();
+        }
+
+        @Override
+        public Set getAllConnectorNodes(ConnectorId connectorId)
+        {
+            return Collections.emptySet();
+        }
+
+        @Override
+        public InternalNode getCurrentNode()
+        {
+            return null;
+        }
+
+        @Override
+        public Set getCoordinators()
+        {
+            return Collections.emptySet();
+        }
+
+        @Override
+        public Set getShuttingDownCoordinator()
+        {
+            return Collections.emptySet();
+        }
+
+        @Override
+        public Set getResourceManagers()
+        {
+            return Collections.emptySet();
+        }
+
+        @Override
+        public Set getCatalogServers()
+        {
+            return Collections.emptySet();
+        }
+
+        @Override
+        public Set getCoordinatorSidecars()
+        {
+            return Collections.emptySet();
+        }
+
+        @Override
+        public AllNodes getAllNodes()
+        {
+            return new AllNodes(
+                    Collections.emptySet(),
+                    Collections.emptySet(),
+                    Collections.emptySet(),
+                    Collections.emptySet(),
+                    Collections.emptySet(),
+                    Collections.emptySet(),
+                    Collections.emptySet());
+        }
+
+        @Override
+        public void refreshNodes()
+        {
+            // No-op for testing
+        }
+
+        @Override
+        public void addNodeChangeListener(Consumer listener)
+        {
+            // No-op for testing
+        }
+
+        @Override
+        public void removeNodeChangeListener(Consumer listener)
+        {
+            // No-op for testing
+        }
+
+        @Override
+        public Optional getNodeLoadMetrics(String nodeIdentifier)
+        {
+            return Optional.empty();
+        }
+    }
+}
diff --git a/presto-main-base/src/test/java/com/facebook/presto/execution/scheduler/clusterOverload/TestCpuMemoryOverloadPolicy.java b/presto-main-base/src/test/java/com/facebook/presto/execution/scheduler/clusterOverload/TestCpuMemoryOverloadPolicy.java
new file mode 100644
index 0000000000000..f3efbad3f549f
--- /dev/null
+++ b/presto-main-base/src/test/java/com/facebook/presto/execution/scheduler/clusterOverload/TestCpuMemoryOverloadPolicy.java
@@ -0,0 +1,254 @@
+/*
+ * Licensed 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 com.facebook.presto.execution.scheduler.clusterOverload;
+
+import com.facebook.presto.client.NodeVersion;
+import com.facebook.presto.execution.ClusterOverloadConfig;
+import com.facebook.presto.metadata.AllNodes;
+import com.facebook.presto.metadata.InternalNode;
+import com.facebook.presto.metadata.InternalNodeManager;
+import com.facebook.presto.spi.ConnectorId;
+import com.facebook.presto.spi.NodeLoadMetrics;
+import com.facebook.presto.spi.NodeState;
+import com.google.common.collect.ImmutableSet;
+import org.testng.annotations.Test;
+
+import java.net.URI;
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.Map;
+import java.util.Optional;
+import java.util.OptionalInt;
+import java.util.Set;
+import java.util.function.Consumer;
+
+import static com.facebook.presto.execution.ClusterOverloadConfig.OVERLOAD_POLICY_CNT_BASED;
+import static com.facebook.presto.execution.ClusterOverloadConfig.OVERLOAD_POLICY_PCT_BASED;
+import static com.facebook.presto.metadata.InternalNode.NodeStatus.ALIVE;
+import static com.facebook.presto.spi.NodePoolType.DEFAULT;
+import static com.facebook.presto.spi.NodeState.ACTIVE;
+import static org.testng.Assert.assertEquals;
+import static org.testng.Assert.assertFalse;
+import static org.testng.Assert.assertTrue;
+
+public class TestCpuMemoryOverloadPolicy
+{
+    private static final NodeVersion TEST_VERSION = new NodeVersion("test");
+    private static final URI TEST_URI = URI.create("http://test.example.com");
+
+    @Test
+    public void testIsClusterOverloadedCountBasedNoOverload()
+    {
+        ClusterOverloadConfig config = new ClusterOverloadConfig()
+                .setAllowedOverloadWorkersCnt(1)
+                .setOverloadPolicyType(OVERLOAD_POLICY_CNT_BASED);
+        CpuMemoryOverloadPolicy policy = new CpuMemoryOverloadPolicy(config);
+
+        // One node is overloaded, but allowed count is 1, so not overloaded
+        InternalNodeManager nodeManager = createNodeManager(ImmutableSet.of(createNode("node1", true, false), createNode("node2", false, false), createNode("node3", false, false)));
+        assertFalse(policy.isClusterOverloaded(nodeManager));
+    }
+
+    @Test
+    public void testIsClusterOverloadedCountBasedOverload()
+    {
+        ClusterOverloadConfig config = new ClusterOverloadConfig()
+                .setAllowedOverloadWorkersCnt(1)
+                .setOverloadPolicyType(OVERLOAD_POLICY_CNT_BASED);
+        CpuMemoryOverloadPolicy policy = new CpuMemoryOverloadPolicy(config);
+
+        // Two nodes are overloaded, but allowed count is 1, so overloaded
+        InternalNodeManager nodeManager = createNodeManager(ImmutableSet.of(createNode("node1", true, false), createNode("node2", false, true), createNode("node3", false, false)));
+        assertTrue(policy.isClusterOverloaded(nodeManager));
+    }
+
+    @Test
+    public void testIsClusterOverloadedPctBasedNoOverload()
+    {
+        ClusterOverloadConfig config = new ClusterOverloadConfig().setAllowedOverloadWorkersPct(0.4).setOverloadPolicyType(OVERLOAD_POLICY_PCT_BASED);
+        CpuMemoryOverloadPolicy policy = new CpuMemoryOverloadPolicy(config);
+
+        // 1 out of 3 nodes (33%) are overloaded, allowed is 40%, so not overloaded
+        InternalNodeManager nodeManager = createNodeManager(ImmutableSet.of(createNode("node1", true, false), createNode("node2", false, false), createNode("node3", false, false)));
+        assertFalse(policy.isClusterOverloaded(nodeManager));
+    }
+
+    @Test
+    public void testIsClusterOverloadedPctBasedOverload()
+    {
+        ClusterOverloadConfig config = new ClusterOverloadConfig().setAllowedOverloadWorkersPct(0.3).setOverloadPolicyType(OVERLOAD_POLICY_PCT_BASED);
+        CpuMemoryOverloadPolicy policy = new CpuMemoryOverloadPolicy(config);
+
+        // 2 out of 5 nodes (40%) are overloaded, allowed is 30%, so overloaded
+        InternalNodeManager nodeManager = createNodeManager(ImmutableSet.of(createNode("node1", true, false), createNode("node2", false, true), createNode("node3", false, false), createNode("node4", false, false), createNode("node5", false, false)));
+        assertTrue(policy.isClusterOverloaded(nodeManager));
+    }
+
+    @Test
+    public void testIsClusterOverloadedBothMetricsOverloaded()
+    {
+        ClusterOverloadConfig config = new ClusterOverloadConfig()
+                .setAllowedOverloadWorkersCnt(0)
+                .setOverloadPolicyType(OVERLOAD_POLICY_CNT_BASED);
+        CpuMemoryOverloadPolicy policy = new CpuMemoryOverloadPolicy(config);
+
+        // Node has both CPU and memory overloaded, should only count as one overloaded node
+        InternalNodeManager nodeManager = createNodeManager(ImmutableSet.of(createNode("node1", true, true)));
+        assertTrue(policy.isClusterOverloaded(nodeManager));
+    }
+
+    @Test
+    public void testIsClusterOverloadedNoNodes()
+    {
+        ClusterOverloadConfig config = new ClusterOverloadConfig()
+                .setAllowedOverloadWorkersCnt(0)
+                .setOverloadPolicyType(OVERLOAD_POLICY_CNT_BASED);
+        CpuMemoryOverloadPolicy policy = new CpuMemoryOverloadPolicy(config);
+
+        // No nodes, should not be overloaded
+        InternalNodeManager nodeManager = createNodeManager(ImmutableSet.of());
+        assertFalse(policy.isClusterOverloaded(nodeManager));
+    }
+
+    @Test
+    public void testGetNameCountBased()
+    {
+        ClusterOverloadConfig config = new ClusterOverloadConfig()
+                .setOverloadPolicyType(OVERLOAD_POLICY_CNT_BASED);
+        CpuMemoryOverloadPolicy policy = new CpuMemoryOverloadPolicy(config);
+
+        assertEquals(policy.getName(), "cpu-memory-overload-cnt");
+    }
+
+    @Test
+    public void testGetNamePctBased()
+    {
+        ClusterOverloadConfig config = new ClusterOverloadConfig()
+                .setOverloadPolicyType(OVERLOAD_POLICY_PCT_BASED);
+        CpuMemoryOverloadPolicy policy = new CpuMemoryOverloadPolicy(config);
+        assertEquals(policy.getName(), "cpu-memory-overload-pct");
+    }
+
+    // Store metrics separately since they're no longer part of InternalNode
+    private static final Map NODE_METRICS = new HashMap<>();
+
+    private static InternalNode createNode(String nodeId, boolean cpuOverload, boolean memoryOverload)
+    {
+        // Store metrics in the map for later retrieval
+        NODE_METRICS.put(nodeId, new NodeLoadMetrics(0.0, 0.0, 0, cpuOverload, memoryOverload));
+        return new InternalNode(
+                nodeId,
+                TEST_URI,
+                OptionalInt.empty(),
+                TEST_VERSION,
+                false,
+                false,
+                false,
+                false,
+                ALIVE,
+                OptionalInt.empty(),
+                DEFAULT);
+    }
+
+    private static InternalNodeManager createNodeManager(Set nodes)
+    {
+        return new InternalNodeManager()
+        {
+            @Override
+            public Set getNodes(NodeState state)
+            {
+                if (state == ACTIVE) {
+                    return nodes;
+                }
+                return ImmutableSet.of();
+            }
+
+            @Override
+            public Set getActiveConnectorNodes(ConnectorId connectorId)
+            {
+                return ImmutableSet.of();
+            }
+
+            @Override
+            public Set getAllConnectorNodes(ConnectorId connectorId)
+            {
+                return Collections.emptySet();
+            }
+
+            @Override
+            public InternalNode getCurrentNode()
+            {
+                return null;
+            }
+
+            @Override
+            public Set getCoordinators()
+            {
+                return ImmutableSet.of();
+            }
+
+            @Override
+            public Set getShuttingDownCoordinator()
+            {
+                return ImmutableSet.of();
+            }
+
+            @Override
+            public Set getResourceManagers()
+            {
+                return ImmutableSet.of();
+            }
+
+            @Override
+            public Set getCatalogServers()
+            {
+                return ImmutableSet.of();
+            }
+
+            @Override
+            public Set getCoordinatorSidecars()
+            {
+                return ImmutableSet.of();
+            }
+
+            @Override
+            public AllNodes getAllNodes()
+            {
+                return new AllNodes(nodes, ImmutableSet.of(), ImmutableSet.of(), ImmutableSet.of(), ImmutableSet.of(), ImmutableSet.of(), ImmutableSet.of());
+            }
+
+            @Override
+            public void refreshNodes()
+            {
+            }
+
+            @Override
+            public void addNodeChangeListener(Consumer listener)
+            {
+            }
+
+            @Override
+            public void removeNodeChangeListener(Consumer listener)
+            {
+            }
+
+            @Override
+            public Optional getNodeLoadMetrics(String nodeIdentifier)
+            {
+                NodeLoadMetrics metrics = NODE_METRICS.get(nodeIdentifier);
+                return metrics != null ? Optional.of(metrics) : Optional.empty();
+            }
+        };
+    }
+}
diff --git a/presto-main-base/src/test/java/com/facebook/presto/metadata/AbstractMockMetadata.java b/presto-main-base/src/test/java/com/facebook/presto/metadata/AbstractMockMetadata.java
index e13cbb5a80c0c..64d33822ed2dc 100644
--- a/presto-main-base/src/test/java/com/facebook/presto/metadata/AbstractMockMetadata.java
+++ b/presto-main-base/src/test/java/com/facebook/presto/metadata/AbstractMockMetadata.java
@@ -35,6 +35,7 @@
 import com.facebook.presto.spi.connector.ConnectorCapabilities;
 import com.facebook.presto.spi.connector.ConnectorOutputMetadata;
 import com.facebook.presto.spi.connector.ConnectorTableVersion;
+import com.facebook.presto.spi.connector.TableFunctionApplicationResult;
 import com.facebook.presto.spi.constraints.TableConstraint;
 import com.facebook.presto.spi.function.SqlFunction;
 import com.facebook.presto.spi.plan.PartitioningHandle;
@@ -685,4 +686,10 @@ public String normalizeIdentifier(Session session, String catalogName, String id
     {
         return identifier.toLowerCase(ENGLISH);
     }
+
+    @Override
+    public Optional> applyTableFunction(Session session, TableFunctionHandle handle)
+    {
+        return Optional.empty();
+    }
 }
diff --git a/presto-native-sidecar-plugin/src/test/java/com/facebook/presto/sidecar/TestConvertApplicableTypeToVariable.java b/presto-main-base/src/test/java/com/facebook/presto/metadata/TestConvertApplicableTypeToVariable.java
similarity index 98%
rename from presto-native-sidecar-plugin/src/test/java/com/facebook/presto/sidecar/TestConvertApplicableTypeToVariable.java
rename to presto-main-base/src/test/java/com/facebook/presto/metadata/TestConvertApplicableTypeToVariable.java
index 93f9ee75f2b82..9e9a8b629124a 100644
--- a/presto-native-sidecar-plugin/src/test/java/com/facebook/presto/sidecar/TestConvertApplicableTypeToVariable.java
+++ b/presto-main-base/src/test/java/com/facebook/presto/metadata/TestConvertApplicableTypeToVariable.java
@@ -11,7 +11,7 @@
  * See the License for the specific language governing permissions and
  * limitations under the License.
  */
-package com.facebook.presto.sidecar;
+package com.facebook.presto.metadata;
 
 import com.facebook.presto.common.type.NamedTypeSignature;
 import com.facebook.presto.common.type.TypeSignature;
@@ -21,8 +21,8 @@
 import java.util.List;
 import java.util.Optional;
 
+import static com.facebook.presto.builtin.tools.WorkerFunctionUtil.convertApplicableTypeToVariable;
 import static com.facebook.presto.common.type.TypeSignature.parseTypeSignature;
-import static com.facebook.presto.sidecar.functionNamespace.NativeFunctionNamespaceManager.convertApplicableTypeToVariable;
 import static org.testng.Assert.assertEquals;
 import static org.testng.Assert.assertTrue;
 
diff --git a/presto-main-base/src/test/java/com/facebook/presto/metadata/TestTableFunctionRegistry.java b/presto-main-base/src/test/java/com/facebook/presto/metadata/TestTableFunctionRegistry.java
index 439d7b35a7c9e..e7670a5af9b5d 100644
--- a/presto-main-base/src/test/java/com/facebook/presto/metadata/TestTableFunctionRegistry.java
+++ b/presto-main-base/src/test/java/com/facebook/presto/metadata/TestTableFunctionRegistry.java
@@ -15,25 +15,18 @@
 
 import com.facebook.presto.Session;
 import com.facebook.presto.spi.ConnectorId;
-import com.facebook.presto.spi.ConnectorSession;
-import com.facebook.presto.spi.connector.ConnectorTransactionHandle;
-import com.facebook.presto.spi.function.table.Argument;
-import com.facebook.presto.spi.function.table.ArgumentSpecification;
-import com.facebook.presto.spi.function.table.ConnectorTableFunction;
-import com.facebook.presto.spi.function.table.ReturnTypeSpecification;
-import com.facebook.presto.spi.function.table.ScalarArgumentSpecification;
-import com.facebook.presto.spi.function.table.TableArgumentSpecification;
-import com.facebook.presto.spi.function.table.TableFunctionAnalysis;
 import com.facebook.presto.spi.security.Identity;
 import com.facebook.presto.sql.tree.QualifiedName;
 import com.google.common.collect.ImmutableList;
 import org.testng.annotations.Test;
 
-import java.util.List;
-import java.util.Map;
 import java.util.Optional;
 
-import static com.facebook.presto.common.type.IntegerType.INTEGER;
+import static com.facebook.presto.connector.tvf.TestingTableFunctions.DuplicateArgumentsTableFunction;
+import static com.facebook.presto.connector.tvf.TestingTableFunctions.MultipleRSTableFunction;
+import static com.facebook.presto.connector.tvf.TestingTableFunctions.NullArgumentsTableFunction;
+import static com.facebook.presto.connector.tvf.TestingTableFunctions.TestConnectorTableFunction;
+import static com.facebook.presto.connector.tvf.TestingTableFunctions.TestConnectorTableFunction2;
 import static com.facebook.presto.testing.TestingSession.testSessionBuilder;
 import static com.facebook.presto.testing.assertions.Assert.assertEquals;
 import static org.testng.Assert.assertFalse;
@@ -44,7 +37,7 @@ public class TestTableFunctionRegistry
 {
     private static final String CATALOG = "test_catalog";
     private static final String USER = "user";
-    private static final String SCHEMA = "test_schema";
+    private static final String SCHEMA = "system";
     private static final Session SESSION = testSessionBuilder()
             .setCatalog(CATALOG)
             .setSchema(SCHEMA)
@@ -86,188 +79,14 @@ public void testTableFunctionRegistry()
 
         // Verify that null arguments table functions cannot be added.
         ex = expectThrows(NullPointerException.class, () -> testFunctionRegistry.addTableFunctions(id, ImmutableList.of(new NullArgumentsTableFunction())));
-        assertTrue(ex.getMessage().contains("table function arguments is null"), ex.getMessage());
+        assertTrue(ex.getMessage().contains("arguments is null"), ex.getMessage());
 
         // Verify that duplicate arguments table functions cannot be added.
         ex = expectThrows(IllegalArgumentException.class, () -> testFunctionRegistry.addTableFunctions(id, ImmutableList.of(new DuplicateArgumentsTableFunction())));
         assertTrue(ex.getMessage().contains("duplicate argument name: a"), ex.getMessage());
 
         // Verify that two row semantic table function arguments functions cannot be added.
-        ex = expectThrows(IllegalArgumentException.class, () -> testFunctionRegistry.addTableFunctions(id, ImmutableList.of(new MultipleRTSTableFunction())));
+        ex = expectThrows(IllegalArgumentException.class, () -> testFunctionRegistry.addTableFunctions(id, ImmutableList.of(new MultipleRSTableFunction())));
         assertTrue(ex.getMessage().contains("more than one table argument with row semantics"), ex.getMessage());
     }
-
-    private static class TestConnectorTableFunction
-            implements ConnectorTableFunction
-    {
-        @Override
-        public String getSchema()
-        {
-            return SCHEMA;
-        }
-
-        @Override
-        public String getName()
-        {
-            return TEST_FUNCTION;
-        }
-
-        @Override
-        public List getArguments()
-        {
-            return ImmutableList.of();
-        }
-
-        @Override
-        public ReturnTypeSpecification getReturnTypeSpecification()
-        {
-            return ReturnTypeSpecification.OnlyPassThrough.ONLY_PASS_THROUGH;
-        }
-
-        @Override
-        public TableFunctionAnalysis analyze(ConnectorSession session, ConnectorTransactionHandle transaction, Map arguments)
-        {
-            return null;
-        }
-    }
-
-    private static class TestConnectorTableFunction2
-            implements ConnectorTableFunction
-    {
-        @Override
-        public String getSchema()
-        {
-            return SCHEMA;
-        }
-
-        @Override
-        public String getName()
-        {
-            return TEST_FUNCTION_2;
-        }
-
-        @Override
-        public List getArguments()
-        {
-            return ImmutableList.of();
-        }
-
-        @Override
-        public ReturnTypeSpecification getReturnTypeSpecification()
-        {
-            return ReturnTypeSpecification.OnlyPassThrough.ONLY_PASS_THROUGH;
-        }
-
-        @Override
-        public TableFunctionAnalysis analyze(ConnectorSession session, ConnectorTransactionHandle transaction, Map arguments)
-        {
-            return null;
-        }
-    }
-
-    private static class NullArgumentsTableFunction
-            implements ConnectorTableFunction
-    {
-        @Override
-        public String getSchema()
-        {
-            return SCHEMA;
-        }
-
-        @Override
-        public String getName()
-        {
-            return TEST_FUNCTION;
-        }
-
-        @Override
-        public List getArguments()
-        {
-            return null;
-        }
-
-        @Override
-        public ReturnTypeSpecification getReturnTypeSpecification()
-        {
-            return ReturnTypeSpecification.OnlyPassThrough.ONLY_PASS_THROUGH;
-        }
-
-        @Override
-        public TableFunctionAnalysis analyze(ConnectorSession session, ConnectorTransactionHandle transaction, Map arguments)
-        {
-            return null;
-        }
-    }
-
-    private static class DuplicateArgumentsTableFunction
-            implements ConnectorTableFunction
-    {
-        @Override
-        public String getSchema()
-        {
-            return SCHEMA;
-        }
-
-        @Override
-        public String getName()
-        {
-            return TEST_FUNCTION;
-        }
-
-        @Override
-        public List getArguments()
-        {
-            ScalarArgumentSpecification s = ScalarArgumentSpecification.builder().name("a").type(INTEGER).build();
-            ScalarArgumentSpecification s2 = ScalarArgumentSpecification.builder().name("a").type(INTEGER).build();
-            return ImmutableList.of(s, s2);
-        }
-
-        @Override
-        public ReturnTypeSpecification getReturnTypeSpecification()
-        {
-            return ReturnTypeSpecification.OnlyPassThrough.ONLY_PASS_THROUGH;
-        }
-
-        @Override
-        public TableFunctionAnalysis analyze(ConnectorSession session, ConnectorTransactionHandle transaction, Map arguments)
-        {
-            return null;
-        }
-    }
-
-    private static class MultipleRTSTableFunction
-            implements ConnectorTableFunction
-    {
-        @Override
-        public String getSchema()
-        {
-            return SCHEMA;
-        }
-
-        @Override
-        public String getName()
-        {
-            return TEST_FUNCTION;
-        }
-
-        @Override
-        public List getArguments()
-        {
-            TableArgumentSpecification t = TableArgumentSpecification.builder().name("t").rowSemantics().build();
-            TableArgumentSpecification t2 = TableArgumentSpecification.builder().name("t2").rowSemantics().build();
-            return ImmutableList.of(t, t2);
-        }
-
-        @Override
-        public ReturnTypeSpecification getReturnTypeSpecification()
-        {
-            return ReturnTypeSpecification.OnlyPassThrough.ONLY_PASS_THROUGH;
-        }
-
-        @Override
-        public TableFunctionAnalysis analyze(ConnectorSession session, ConnectorTransactionHandle transaction, Map arguments)
-        {
-            return null;
-        }
-    }
 }
diff --git a/presto-main-base/src/test/java/com/facebook/presto/operator/TestAnnotationEngineForAggregates.java b/presto-main-base/src/test/java/com/facebook/presto/operator/TestAnnotationEngineForAggregates.java
index 5927f1fc65ee6..220b91ba548d8 100644
--- a/presto-main-base/src/test/java/com/facebook/presto/operator/TestAnnotationEngineForAggregates.java
+++ b/presto-main-base/src/test/java/com/facebook/presto/operator/TestAnnotationEngineForAggregates.java
@@ -43,6 +43,7 @@
 import com.facebook.presto.spi.function.Description;
 import com.facebook.presto.spi.function.FunctionKind;
 import com.facebook.presto.spi.function.InputFunction;
+import com.facebook.presto.spi.function.LiteralParameter;
 import com.facebook.presto.spi.function.LiteralParameters;
 import com.facebook.presto.spi.function.LongVariableConstraint;
 import com.facebook.presto.spi.function.OperatorDependency;
@@ -53,7 +54,6 @@
 import com.facebook.presto.spi.function.TypeParameterSpecialization;
 import com.facebook.presto.spi.function.aggregation.AggregationMetadata;
 import com.facebook.presto.type.Constraint;
-import com.facebook.presto.type.LiteralParameter;
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableSet;
 import io.airlift.slice.Slice;
diff --git a/presto-main-base/src/test/java/com/facebook/presto/operator/TestAnnotationEngineForScalars.java b/presto-main-base/src/test/java/com/facebook/presto/operator/TestAnnotationEngineForScalars.java
index 80ab4502de2e6..787a2a5f8da27 100644
--- a/presto-main-base/src/test/java/com/facebook/presto/operator/TestAnnotationEngineForScalars.java
+++ b/presto-main-base/src/test/java/com/facebook/presto/operator/TestAnnotationEngineForScalars.java
@@ -30,13 +30,13 @@
 import com.facebook.presto.spi.function.Description;
 import com.facebook.presto.spi.function.FunctionKind;
 import com.facebook.presto.spi.function.IsNull;
+import com.facebook.presto.spi.function.LiteralParameter;
 import com.facebook.presto.spi.function.LiteralParameters;
 import com.facebook.presto.spi.function.ScalarFunction;
 import com.facebook.presto.spi.function.Signature;
 import com.facebook.presto.spi.function.SqlNullable;
 import com.facebook.presto.spi.function.SqlType;
 import com.facebook.presto.spi.function.TypeParameter;
-import com.facebook.presto.type.LiteralParameter;
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableSet;
 import io.airlift.slice.Slice;
diff --git a/presto-main-base/src/test/java/com/facebook/presto/operator/scalar/AbstractTestFunctions.java b/presto-main-base/src/test/java/com/facebook/presto/operator/scalar/AbstractTestFunctions.java
index d31ae0b09eb02..0644a6ebbdb68 100644
--- a/presto-main-base/src/test/java/com/facebook/presto/operator/scalar/AbstractTestFunctions.java
+++ b/presto-main-base/src/test/java/com/facebook/presto/operator/scalar/AbstractTestFunctions.java
@@ -59,6 +59,7 @@ public abstract class AbstractTestFunctions
     private final FeaturesConfig featuresConfig;
     private final FunctionsConfig functionsConfig;
     protected FunctionAssertions functionAssertions;
+    private final boolean loadInlinedSqlInvokedFunctionsPlugin;
 
     protected AbstractTestFunctions()
     {
@@ -81,18 +82,23 @@ protected AbstractTestFunctions(FunctionsConfig functionsConfig)
     }
 
     protected AbstractTestFunctions(Session session, FeaturesConfig featuresConfig, FunctionsConfig functionsConfig)
+    {
+        this(session, featuresConfig, functionsConfig, true);
+    }
+    protected AbstractTestFunctions(Session session, FeaturesConfig featuresConfig, FunctionsConfig functionsConfig, boolean loadInlinedSqlInvokedFunctionsPlugin)
     {
         this.session = requireNonNull(session, "session is null");
         this.featuresConfig = requireNonNull(featuresConfig, "featuresConfig is null");
         this.functionsConfig = requireNonNull(functionsConfig, "config is null")
                 .setLegacyLogFunction(true)
                 .setUseNewNanDefinition(true);
+        this.loadInlinedSqlInvokedFunctionsPlugin = loadInlinedSqlInvokedFunctionsPlugin;
     }
 
     @BeforeClass
     public final void initTestFunctions()
     {
-        functionAssertions = new FunctionAssertions(session, featuresConfig, functionsConfig, false);
+        functionAssertions = new FunctionAssertions(session, featuresConfig, functionsConfig, false, loadInlinedSqlInvokedFunctionsPlugin);
     }
 
     @AfterClass(alwaysRun = true)
diff --git a/presto-main-base/src/test/java/com/facebook/presto/operator/scalar/FunctionAssertions.java b/presto-main-base/src/test/java/com/facebook/presto/operator/scalar/FunctionAssertions.java
index 17ba3c70823db..d46bb9a0d0bd2 100644
--- a/presto-main-base/src/test/java/com/facebook/presto/operator/scalar/FunctionAssertions.java
+++ b/presto-main-base/src/test/java/com/facebook/presto/operator/scalar/FunctionAssertions.java
@@ -41,6 +41,7 @@
 import com.facebook.presto.operator.project.CursorProcessor;
 import com.facebook.presto.operator.project.PageProcessor;
 import com.facebook.presto.operator.project.PageProjectionWithOutputs;
+import com.facebook.presto.scalar.sql.SqlInvokedFunctionsPlugin;
 import com.facebook.presto.spi.ColumnHandle;
 import com.facebook.presto.spi.ConnectorId;
 import com.facebook.presto.spi.ConnectorPageSource;
@@ -221,20 +222,25 @@ public FunctionAssertions()
 
     public FunctionAssertions(Session session)
     {
-        this(session, new FeaturesConfig(), new FunctionsConfig(), false);
+        this(session, new FeaturesConfig(), new FunctionsConfig(), false, true);
     }
 
     public FunctionAssertions(Session session, FeaturesConfig featuresConfig)
     {
-        this(session, featuresConfig, new FunctionsConfig(), false);
+        this(session, featuresConfig, new FunctionsConfig(), false, true);
     }
 
     public FunctionAssertions(Session session, FunctionsConfig functionsConfig)
     {
-        this(session, new FeaturesConfig(), functionsConfig, false);
+        this(session, new FeaturesConfig(), functionsConfig, false, true);
     }
 
     public FunctionAssertions(Session session, FeaturesConfig featuresConfig, FunctionsConfig functionsConfig, boolean refreshSession)
+    {
+        this(session, featuresConfig, functionsConfig, refreshSession, true);
+    }
+
+    public FunctionAssertions(Session session, FeaturesConfig featuresConfig, FunctionsConfig functionsConfig, boolean refreshSession, boolean loadInlinedSqlInvokedFunctionsPlugin)
     {
         requireNonNull(session, "session is null");
         runner = new LocalQueryRunner(session, featuresConfig, functionsConfig);
@@ -244,6 +250,9 @@ public FunctionAssertions(Session session, FeaturesConfig featuresConfig, Functi
         else {
             this.session = session;
         }
+        if (loadInlinedSqlInvokedFunctionsPlugin) {
+            runner.installPlugin(new SqlInvokedFunctionsPlugin());
+        }
         metadata = runner.getMetadata();
         compiler = runner.getExpressionCompiler();
     }
@@ -269,6 +278,12 @@ public FunctionAssertions addScalarFunctions(Class clazz)
         return this;
     }
 
+    public FunctionAssertions addConnectorFunctions(List functionInfos, String namespace)
+    {
+        metadata.registerConnectorFunctions(namespace, functionInfos);
+        return this;
+    }
+
     public void assertFunction(String projection, Type expectedType, Object expected)
     {
         if (expected instanceof Slice) {
diff --git a/presto-main-base/src/test/java/com/facebook/presto/operator/scalar/TestArraySortByKeyFunctions.java b/presto-main-base/src/test/java/com/facebook/presto/operator/scalar/TestArraySortByKeyFunctions.java
new file mode 100644
index 0000000000000..dee4f22e42a96
--- /dev/null
+++ b/presto-main-base/src/test/java/com/facebook/presto/operator/scalar/TestArraySortByKeyFunctions.java
@@ -0,0 +1,309 @@
+/*
+ * Licensed 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 com.facebook.presto.operator.scalar;
+
+import com.facebook.presto.common.type.ArrayType;
+import com.facebook.presto.common.type.RowType;
+import com.google.common.collect.ImmutableList;
+import org.testng.annotations.Test;
+
+import static com.facebook.presto.common.type.BigintType.BIGINT;
+import static com.facebook.presto.common.type.BooleanType.BOOLEAN;
+import static com.facebook.presto.common.type.DoubleType.DOUBLE;
+import static com.facebook.presto.common.type.IntegerType.INTEGER;
+import static com.facebook.presto.common.type.UnknownType.UNKNOWN;
+import static com.facebook.presto.common.type.VarcharType.createVarcharType;
+import static java.util.Arrays.asList;
+
+public class TestArraySortByKeyFunctions
+        extends AbstractTestFunctions
+{
+    @Test
+    public void testBasic()
+    {
+        // Test array_sort
+        assertFunction(
+                "array_sort(ARRAY['pear', 'apple', 'banana', 'kiwi'], x -> length(x))",
+                new ArrayType(createVarcharType(6)),
+                asList("pear", "kiwi", "apple", "banana"));
+
+        assertFunction(
+                "array_sort(ARRAY['pear', 'apple', 'banana', 'kiwi'], x -> substr(x, length(x), 1))",
+                new ArrayType(createVarcharType(6)),
+                asList("banana", "apple", "kiwi", "pear"));
+
+        // Test array_sort_desc
+        assertFunction(
+                "array_sort_desc(ARRAY['pear', 'apple', 'banana', 'kiwi'], x -> length(x))",
+                new ArrayType(createVarcharType(6)),
+                asList("banana", "apple", "pear", "kiwi"));
+
+        assertFunction(
+                "array_sort_desc(ARRAY['pear', 'apple', 'banana', 'kiwi'], x -> substr(x, length(x), 1))",
+                new ArrayType(createVarcharType(6)),
+                asList("pear", "kiwi", "apple", "banana"));
+    }
+
+    @Test
+    public void testNulls()
+    {
+        // Test array_sort
+        assertFunction(
+                "array_sort(ARRAY['apple', NULL, 'banana', NULL], x -> length(x))",
+                new ArrayType(createVarcharType(6)),
+                asList("apple", "banana", null, null));
+
+        assertFunction(
+                "array_sort(ARRAY['apple', 'banana', 'pear'], x -> IF(x = 'banana', NULL, length(x)))",
+                new ArrayType(createVarcharType(6)),
+                asList("pear", "apple", "banana"));
+
+        assertFunction(
+                "array_sort(ARRAY['apple', NULL, 'banana', 'pear', NULL], x -> length(x))",
+                new ArrayType(createVarcharType(6)),
+                asList("pear", "apple", "banana", null, null));
+
+        // Test array_sort_desc
+        assertFunction(
+                "array_sort_desc(ARRAY['apple', NULL, 'banana', NULL], x -> length(x))",
+                new ArrayType(createVarcharType(6)),
+                asList("banana", "apple", null, null));
+
+        assertFunction(
+                "array_sort_desc(ARRAY['apple', 'banana', 'pear'], x -> IF(x = 'banana', NULL, length(x)))",
+                new ArrayType(createVarcharType(6)),
+                asList("apple", "pear", "banana"));
+
+        assertFunction(
+                "array_sort_desc(ARRAY['apple', NULL, 'banana', 'pear', NULL], x -> length(x))",
+                new ArrayType(createVarcharType(6)),
+                asList("banana", "apple", "pear", null, null));
+    }
+
+    @Test
+    public void testSpecialDoubleValues()
+    {
+        // Test array_sort
+        assertFunction(
+                "array_sort(ARRAY[CAST(0.0 AS DOUBLE), CAST('NaN' AS DOUBLE), CAST('Infinity' AS DOUBLE), CAST('-Infinity' AS DOUBLE)], x -> x)",
+                new ArrayType(DOUBLE),
+                asList(Double.NEGATIVE_INFINITY, 0.0, Double.POSITIVE_INFINITY, Double.NaN));
+
+        // Test array_sort_desc
+        assertFunction(
+                "array_sort_desc(ARRAY[CAST(0.0 AS DOUBLE), CAST('NaN' AS DOUBLE), CAST('Infinity' AS DOUBLE), CAST('-Infinity' AS DOUBLE)], x -> x)",
+                new ArrayType(DOUBLE),
+                asList(Double.NaN, Double.POSITIVE_INFINITY, 0.0, Double.NEGATIVE_INFINITY));
+    }
+
+    @Test
+    public void testNumericKeys()
+    {
+        // Test array_sort
+        assertFunction(
+                "array_sort(ARRAY[5, 20, 3, 9, 100], x -> x)",
+                new ArrayType(INTEGER),
+                asList(3, 5, 9, 20, 100));
+
+        assertFunction(
+                "array_sort(ARRAY[CAST(5000000000 AS BIGINT), CAST(20000000000 AS BIGINT), CAST(3000000000 AS BIGINT), CAST(9000000000 AS BIGINT), CAST(100000000000 AS BIGINT)], x -> x)",
+                new ArrayType(BIGINT),
+                asList(3000000000L, 5000000000L, 9000000000L, 20000000000L, 100000000000L));
+
+        assertFunction(
+                "array_sort(ARRAY[CAST(5.5 AS DOUBLE), CAST(20.1 AS DOUBLE), CAST(3.9 AS DOUBLE), CAST(9.0 AS DOUBLE), CAST(100.0 AS DOUBLE)], x -> x)",
+                new ArrayType(DOUBLE),
+                asList(3.9, 5.5, 9.0, 20.1, 100.0));
+
+        assertFunction(
+                "array_sort(ARRAY[5, 20, 3, 9, 100], x -> x % 10)",
+                new ArrayType(INTEGER),
+                asList(20, 100, 3, 5, 9));
+
+        assertFunction(
+                "array_sort(ARRAY[CAST(5000000000 AS BIGINT), CAST(20000000000 AS BIGINT), CAST(3000000000 AS BIGINT), CAST(9000000000 AS BIGINT), CAST(100000000000 AS BIGINT)], x -> x % CAST(10000000000 AS BIGINT))",
+                new ArrayType(BIGINT),
+                asList(20000000000L, 100000000000L, 3000000000L, 5000000000L, 9000000000L));
+
+        // Test array_sort_desc
+        assertFunction(
+                "array_sort_desc(ARRAY[5, 20, 3, 9, 100], x -> x)",
+                new ArrayType(INTEGER),
+                asList(100, 20, 9, 5, 3));
+
+        assertFunction(
+                "array_sort_desc(ARRAY[CAST(5000000000 AS BIGINT), CAST(20000000000 AS BIGINT), CAST(3000000000 AS BIGINT), CAST(9000000000 AS BIGINT), CAST(100000000000 AS BIGINT)], x -> x)",
+                new ArrayType(BIGINT),
+                asList(100000000000L, 20000000000L, 9000000000L, 5000000000L, 3000000000L));
+
+        assertFunction(
+                "array_sort_desc(ARRAY[CAST(5.5 AS DOUBLE), CAST(20.1 AS DOUBLE), CAST(3.9 AS DOUBLE), CAST(9.0 AS DOUBLE), CAST(100.0 AS DOUBLE)], x -> x)",
+                new ArrayType(DOUBLE),
+                asList(100.0, 20.1, 9.0, 5.5, 3.9));
+
+        assertFunction(
+                "array_sort_desc(ARRAY[5, 20, 3, 9, 100], x -> x % 10)",
+                new ArrayType(INTEGER),
+                asList(9, 5, 3, 20, 100));
+
+        assertFunction(
+                "array_sort_desc(ARRAY[CAST(5000000000 AS BIGINT), CAST(20000000000 AS BIGINT), CAST(3000000000 AS BIGINT), CAST(9000000000 AS BIGINT), CAST(100000000000 AS BIGINT)], x -> x % CAST(10000000000 AS BIGINT))",
+                new ArrayType(BIGINT),
+                asList(9000000000L, 5000000000L, 3000000000L, 20000000000L, 100000000000L));
+    }
+
+    @Test
+    public void testBooleanKeys()
+    {
+        // Test array_sort
+        assertFunction(
+                "array_sort(ARRAY[true, false, true, false], x -> x)",
+                new ArrayType(BOOLEAN),
+                asList(false, false, true, true));
+
+        assertFunction(
+                "array_sort(ARRAY[true, false, true, false], x -> NOT x)",
+                new ArrayType(BOOLEAN),
+                asList(true, true, false, false));
+
+        // Test array_sort_desc
+        assertFunction(
+                "array_sort_desc(ARRAY[true, false, true, false], x -> x)",
+                new ArrayType(BOOLEAN),
+                asList(true, true, false, false));
+
+        assertFunction(
+                "array_sort_desc(ARRAY[true, false, true, false], x -> NOT x)",
+                new ArrayType(BOOLEAN),
+                asList(false, false, true, true));
+    }
+
+    @Test
+    public void testComplexTypes()
+    {
+        // Test array_sort
+        assertFunction(
+                "array_sort(ARRAY[ARRAY[1, 2, 3], ARRAY[4, 5], ARRAY[6, 7, 8, 9]], x -> cardinality(x))",
+                new ArrayType(new ArrayType(INTEGER)),
+                asList(asList(4, 5), asList(1, 2, 3), asList(6, 7, 8, 9)));
+
+        assertFunction(
+                "array_sort(ARRAY[ROW('a', 3), ROW('b', 1), ROW('c', 2)], x -> x[2])",
+                new ArrayType(RowType.anonymous(ImmutableList.of(createVarcharType(1), INTEGER))),
+                asList(asList("b", 1), asList("c", 2), asList("a", 3)));
+
+        assertFunction(
+                "array_sort(ARRAY[ROW('a', CAST(3000000000 AS BIGINT)), ROW('b', CAST(1000000000 AS BIGINT)), ROW('c', CAST(2000000000 AS BIGINT))], x -> x[2])",
+                new ArrayType(RowType.anonymous(ImmutableList.of(createVarcharType(1), BIGINT))),
+                asList(asList("b", 1000000000L), asList("c", 2000000000L), asList("a", 3000000000L)));
+
+        // Test array_sort_desc
+        assertFunction(
+                "array_sort_desc(ARRAY[ARRAY[1, 2, 3], ARRAY[4, 5], ARRAY[6, 7, 8, 9]], x -> cardinality(x))",
+                new ArrayType(new ArrayType(INTEGER)),
+                asList(asList(6, 7, 8, 9), asList(1, 2, 3), asList(4, 5)));
+
+        assertFunction(
+                "array_sort_desc(ARRAY[ROW('a', 3), ROW('b', 1), ROW('c', 2)], x -> x[2])",
+                new ArrayType(RowType.anonymous(ImmutableList.of(createVarcharType(1), INTEGER))),
+                asList(asList("a", 3), asList("c", 2), asList("b", 1)));
+
+        assertFunction(
+                "array_sort_desc(ARRAY[ROW('a', CAST(3000000000 AS BIGINT)), ROW('b', CAST(1000000000 AS BIGINT)), ROW('c', CAST(2000000000 AS BIGINT))], x -> x[2])",
+                new ArrayType(RowType.anonymous(ImmutableList.of(createVarcharType(1), BIGINT))),
+                asList(asList("a", 3000000000L), asList("c", 2000000000L), asList("b", 1000000000L)));
+    }
+
+    @Test
+    public void testEdgeCases()
+    {
+        // Test array_sort
+        assertFunction(
+                "array_sort(ARRAY[], x -> x)",
+                new ArrayType(UNKNOWN),
+                ImmutableList.of());
+
+        assertFunction(
+                "array_sort(ARRAY[5], x -> x)",
+                new ArrayType(INTEGER),
+                asList(5));
+
+        assertFunction(
+                "array_sort(ARRAY[NULL, NULL, NULL], x -> x)",
+                new ArrayType(UNKNOWN),
+                asList(null, null, null));
+
+        // Test array_sort_desc
+        assertFunction(
+                "array_sort_desc(ARRAY[], x -> x)",
+                new ArrayType(UNKNOWN),
+                ImmutableList.of());
+
+        assertFunction(
+                "array_sort_desc(ARRAY[5], x -> x)",
+                new ArrayType(INTEGER),
+                asList(5));
+
+        assertFunction(
+                "array_sort_desc(ARRAY[NULL, NULL, NULL], x -> x)",
+                new ArrayType(UNKNOWN),
+                asList(null, null, null));
+    }
+
+    @Test
+    public void testTypeCoercion()
+    {
+        // Test array_sort
+        assertFunction(
+                "array_sort(ARRAY[5, 20, 3, 9, 100], x -> x + CAST(0.5 AS DOUBLE))",
+                new ArrayType(INTEGER),
+                asList(3, 5, 9, 20, 100));
+
+        assertFunction(
+                "array_sort(ARRAY[5, 20, 3, 9, 100], x -> x * CAST(1000000000 AS BIGINT))",
+                new ArrayType(INTEGER),
+                asList(3, 5, 9, 20, 100));
+
+        assertFunction(
+                "array_sort(ARRAY['5', '20', '3', '9', '100'], x -> cast(x as integer))",
+                new ArrayType(createVarcharType(3)),
+                asList("3", "5", "9", "20", "100"));
+
+        assertFunction(
+                "array_sort(ARRAY['5000000000', '20000000000', '3000000000', '9000000000', '100000000000'], x -> cast(x as bigint))",
+                new ArrayType(createVarcharType(12)),
+                asList("3000000000", "5000000000", "9000000000", "20000000000", "100000000000"));
+
+        // Test array_sort_desc
+        assertFunction(
+                "array_sort_desc(ARRAY[5, 20, 3, 9, 100], x -> x + CAST(0.5 AS DOUBLE))",
+                new ArrayType(INTEGER),
+                asList(100, 20, 9, 5, 3));
+
+        assertFunction(
+                "array_sort_desc(ARRAY[5, 20, 3, 9, 100], x -> x * CAST(1000000000 AS BIGINT))",
+                new ArrayType(INTEGER),
+                asList(100, 20, 9, 5, 3));
+
+        assertFunction(
+                "array_sort_desc(ARRAY['5', '20', '3', '9', '100'], x -> cast(x as integer))",
+                new ArrayType(createVarcharType(3)),
+                asList("100", "20", "9", "5", "3"));
+
+        assertFunction(
+                "array_sort_desc(ARRAY['5000000000', '20000000000', '3000000000', '9000000000', '100000000000'], x -> cast(x as bigint))",
+                new ArrayType(createVarcharType(12)),
+                asList("100000000000", "20000000000", "9000000000", "5000000000", "3000000000"));
+    }
+}
diff --git a/presto-main-base/src/test/java/com/facebook/presto/operator/scalar/TestDateTimeFunctions.java b/presto-main-base/src/test/java/com/facebook/presto/operator/scalar/TestDateTimeFunctions.java
index 1a939a01b3461..cf45e5643f9cb 100644
--- a/presto-main-base/src/test/java/com/facebook/presto/operator/scalar/TestDateTimeFunctions.java
+++ b/presto-main-base/src/test/java/com/facebook/presto/operator/scalar/TestDateTimeFunctions.java
@@ -15,12 +15,10 @@
 package com.facebook.presto.operator.scalar;
 
 import com.facebook.presto.Session;
-import com.facebook.presto.common.type.TimeType;
 import com.facebook.presto.common.type.TimestampType;
 import org.joda.time.DateTime;
 import org.testng.annotations.Test;
 
-import static com.facebook.presto.common.type.TimeWithTimeZoneType.TIME_WITH_TIME_ZONE;
 import static com.facebook.presto.common.type.TimestampWithTimeZoneType.TIMESTAMP_WITH_TIME_ZONE;
 import static com.facebook.presto.common.type.VarcharType.createVarcharType;
 
@@ -46,31 +44,6 @@ public void testFormatDateCannotImplicitlyAddTimeZoneToTimestampLiteral()
                 "format_datetime for TIMESTAMP type, cannot use 'Z' nor 'z' in format, as this type does not contain TZ information");
     }
 
-    @Test
-    public void testLocalTime()
-    {
-        Session localSession = Session.builder(session)
-                .setStartTime(new DateTime(2017, 3, 1, 14, 30, 0, 0, DATE_TIME_ZONE).getMillis())
-                .build();
-        try (FunctionAssertions localAssertion = new FunctionAssertions(localSession)) {
-            localAssertion.assertFunctionString("LOCALTIME", TimeType.TIME, "14:30:00.000");
-        }
-    }
-
-    @Test
-    public void testCurrentTime()
-    {
-        Session localSession = Session.builder(session)
-                // we use Asia/Kathmandu here to test the difference in semantic change of current_time
-                // between legacy and non-legacy timestamp
-                .setTimeZoneKey(KATHMANDU_ZONE_KEY)
-                .setStartTime(new DateTime(2017, 3, 1, 15, 45, 0, 0, KATHMANDU_ZONE).getMillis())
-                .build();
-        try (FunctionAssertions localAssertion = new FunctionAssertions(localSession)) {
-            localAssertion.assertFunctionString("CURRENT_TIME", TIME_WITH_TIME_ZONE, "15:45:00.000 Asia/Kathmandu");
-        }
-    }
-
     @Test
     public void testLocalTimestamp()
     {
diff --git a/presto-main-base/src/test/java/com/facebook/presto/operator/scalar/TestDateTimeFunctionsBase.java b/presto-main-base/src/test/java/com/facebook/presto/operator/scalar/TestDateTimeFunctionsBase.java
index 7fa1263fcd5e1..733108e7d8dcc 100644
--- a/presto-main-base/src/test/java/com/facebook/presto/operator/scalar/TestDateTimeFunctionsBase.java
+++ b/presto-main-base/src/test/java/com/facebook/presto/operator/scalar/TestDateTimeFunctionsBase.java
@@ -186,6 +186,30 @@ private static long epochDaysInZone(TimeZoneKey timeZoneKey, long instant)
         return LocalDate.from(Instant.ofEpochMilli(instant).atZone(ZoneId.of(timeZoneKey.getId()))).toEpochDay();
     }
 
+    @Test
+    public void testLocalTime()
+    {
+        Session localSession = Session.builder(session)
+                .setStartTime(new DateTime(2017, 3, 1, 14, 30, 0, 0, DATE_TIME_ZONE).getMillis())
+                .build();
+        try (FunctionAssertions localAssertion = new FunctionAssertions(localSession)) {
+            localAssertion.assertFunctionString("LOCALTIME", TimeType.TIME, "14:30:00.000");
+        }
+    }
+
+    @Test
+    public void testCurrentTime()
+    {
+        Session localSession = Session.builder(session)
+                // we use Asia/Kathmandu here, as it has different zone offset on 2017-03-01 and on 1970-01-01
+                .setTimeZoneKey(KATHMANDU_ZONE_KEY)
+                .setStartTime(new DateTime(2017, 3, 1, 15, 45, 0, 0, KATHMANDU_ZONE).getMillis())
+                .build();
+        try (FunctionAssertions localAssertion = new FunctionAssertions(localSession)) {
+            localAssertion.assertFunctionString("CURRENT_TIME", TIME_WITH_TIME_ZONE, "15:45:00.000 Asia/Kathmandu");
+        }
+    }
+
     @Test
     public void testFromUnixTime()
     {
diff --git a/presto-main-base/src/test/java/com/facebook/presto/operator/scalar/TestDateTimeFunctionsLegacy.java b/presto-main-base/src/test/java/com/facebook/presto/operator/scalar/TestDateTimeFunctionsLegacy.java
index 97527e6b00d10..2293dd0ceb865 100644
--- a/presto-main-base/src/test/java/com/facebook/presto/operator/scalar/TestDateTimeFunctionsLegacy.java
+++ b/presto-main-base/src/test/java/com/facebook/presto/operator/scalar/TestDateTimeFunctionsLegacy.java
@@ -15,12 +15,10 @@
 package com.facebook.presto.operator.scalar;
 
 import com.facebook.presto.Session;
-import com.facebook.presto.common.type.TimeType;
 import com.facebook.presto.common.type.TimestampType;
 import org.joda.time.DateTime;
 import org.testng.annotations.Test;
 
-import static com.facebook.presto.common.type.TimeWithTimeZoneType.TIME_WITH_TIME_ZONE;
 import static com.facebook.presto.common.type.TimestampWithTimeZoneType.TIMESTAMP_WITH_TIME_ZONE;
 import static com.facebook.presto.common.type.VarcharType.VARCHAR;
 import static com.facebook.presto.common.type.VarcharType.createVarcharType;
@@ -45,31 +43,6 @@ public void testFormatDateCanImplicitlyAddTimeZoneToTimestampLiteral()
         assertFunction("format_datetime(" + TIMESTAMP_LITERAL + ", 'YYYY/MM/dd HH:mm ZZZZ')", VARCHAR, "2001/08/22 03:04 " + DATE_TIME_ZONE.getID());
     }
 
-    @Test
-    public void testLocalTime()
-    {
-        Session localSession = Session.builder(session)
-                .setStartTime(new DateTime(2017, 3, 1, 14, 30, 0, 0, DATE_TIME_ZONE).getMillis())
-                .build();
-        try (FunctionAssertions localAssertion = new FunctionAssertions(localSession)) {
-            localAssertion.assertFunctionString("LOCALTIME", TimeType.TIME, "13:30:00.000");
-        }
-    }
-
-    @Test
-    public void testCurrentTime()
-    {
-        Session localSession = Session.builder(session)
-                // we use Asia/Kathmandu here to test the difference in semantic change of current_time
-                // between legacy and non-legacy timestamp
-                .setTimeZoneKey(KATHMANDU_ZONE_KEY)
-                .setStartTime(new DateTime(2017, 3, 1, 15, 45, 0, 0, KATHMANDU_ZONE).getMillis())
-                .build();
-        try (FunctionAssertions localAssertion = new FunctionAssertions(localSession)) {
-            localAssertion.assertFunctionString("CURRENT_TIME", TIME_WITH_TIME_ZONE, "15:30:00.000 Asia/Kathmandu");
-        }
-    }
-
     @Test
     public void testLocalTimestamp()
     {
diff --git a/presto-main-base/src/test/java/com/facebook/presto/operator/scalar/TestStringFunctions.java b/presto-main-base/src/test/java/com/facebook/presto/operator/scalar/TestStringFunctions.java
index 9197fae5aa811..0023402dd9e91 100644
--- a/presto-main-base/src/test/java/com/facebook/presto/operator/scalar/TestStringFunctions.java
+++ b/presto-main-base/src/test/java/com/facebook/presto/operator/scalar/TestStringFunctions.java
@@ -18,10 +18,10 @@
 import com.facebook.presto.common.type.SqlVarbinary;
 import com.facebook.presto.common.type.StandardTypes;
 import com.facebook.presto.spi.function.Description;
+import com.facebook.presto.spi.function.LiteralParameter;
 import com.facebook.presto.spi.function.LiteralParameters;
 import com.facebook.presto.spi.function.ScalarFunction;
 import com.facebook.presto.spi.function.SqlType;
-import com.facebook.presto.type.LiteralParameter;
 import com.google.common.base.Joiner;
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableMap;
diff --git a/presto-main-base/src/test/java/com/facebook/presto/server/TestInternalCommunicationConfig.java b/presto-main-base/src/test/java/com/facebook/presto/server/TestInternalCommunicationConfig.java
index ba6faf2a97c53..de88e6b3fc877 100644
--- a/presto-main-base/src/test/java/com/facebook/presto/server/TestInternalCommunicationConfig.java
+++ b/presto-main-base/src/test/java/com/facebook/presto/server/TestInternalCommunicationConfig.java
@@ -52,7 +52,9 @@ public void testDefaults()
                 .setSharedSecret(null)
                 .setTaskUpdateRequestThriftSerdeEnabled(false)
                 .setTaskInfoResponseThriftSerdeEnabled(false)
-                .setInternalJwtEnabled(false));
+                .setInternalJwtEnabled(false)
+                .setNodeStatsRefreshIntervalMillis(1_000)
+                .setNodeDiscoveryPollingIntervalMillis(5_000));
     }
 
     @Test
@@ -80,6 +82,8 @@ public void testExplicitPropertyMappings()
                 .put("internal-communication.jwt.enabled", "true")
                 .put("experimental.internal-communication.task-update-request-thrift-serde-enabled", "true")
                 .put("experimental.internal-communication.task-info-response-thrift-serde-enabled", "true")
+                .put("internal-communication.node-stats-refresh-interval-millis", "2000")
+                .put("internal-communication.node-discovery-polling-interval-millis", "3000")
                 .build();
 
         InternalCommunicationConfig expected = new InternalCommunicationConfig()
@@ -103,7 +107,9 @@ public void testExplicitPropertyMappings()
                 .setSharedSecret("secret")
                 .setInternalJwtEnabled(true)
                 .setTaskUpdateRequestThriftSerdeEnabled(true)
-                .setTaskInfoResponseThriftSerdeEnabled(true);
+                .setTaskInfoResponseThriftSerdeEnabled(true)
+                .setNodeStatsRefreshIntervalMillis(2000)
+                .setNodeDiscoveryPollingIntervalMillis(3000);
 
         assertFullMapping(properties, expected);
     }
diff --git a/presto-main-base/src/test/java/com/facebook/presto/server/TestQueryStateInfo.java b/presto-main-base/src/test/java/com/facebook/presto/server/TestQueryStateInfo.java
index 09230829b8e23..c4bd91d43640f 100644
--- a/presto-main-base/src/test/java/com/facebook/presto/server/TestQueryStateInfo.java
+++ b/presto-main-base/src/test/java/com/facebook/presto/server/TestQueryStateInfo.java
@@ -17,11 +17,15 @@
 import com.facebook.airlift.units.Duration;
 import com.facebook.presto.common.RuntimeStats;
 import com.facebook.presto.cost.StatsAndCosts;
+import com.facebook.presto.execution.ClusterOverloadConfig;
 import com.facebook.presto.execution.QueryInfo;
 import com.facebook.presto.execution.QueryState;
 import com.facebook.presto.execution.QueryStats;
 import com.facebook.presto.execution.resourceGroups.InternalResourceGroup;
+import com.facebook.presto.execution.scheduler.clusterOverload.ClusterOverloadPolicy;
+import com.facebook.presto.execution.scheduler.clusterOverload.ClusterResourceChecker;
 import com.facebook.presto.metadata.InMemoryNodeManager;
+import com.facebook.presto.metadata.InternalNodeManager;
 import com.facebook.presto.spi.PrestoWarning;
 import com.facebook.presto.spi.QueryId;
 import com.facebook.presto.spi.WarningCode;
@@ -57,7 +61,7 @@ public class TestQueryStateInfo
     @Test
     public void testQueryStateInfo()
     {
-        InternalResourceGroup.RootInternalResourceGroup root = new InternalResourceGroup.RootInternalResourceGroup("root", (group, export) -> {}, directExecutor(), ignored -> Optional.empty(), rg -> false, new InMemoryNodeManager());
+        InternalResourceGroup.RootInternalResourceGroup root = new InternalResourceGroup.RootInternalResourceGroup("root", (group, export) -> {}, directExecutor(), ignored -> Optional.empty(), rg -> false, new InMemoryNodeManager(), createClusterResourceChecker());
         root.setSoftMemoryLimit(new DataSize(1, MEGABYTE));
         root.setMaxQueuedQueries(40);
         root.setHardConcurrencyLimit(0);
@@ -317,4 +321,29 @@ private QueryInfo createQueryInfo(String queryId, ResourceGroupId resourceGroupI
                 ImmutableMap.of(),
                 Optional.empty());
     }
+
+    private ClusterResourceChecker createClusterResourceChecker()
+    {
+        // Create a mock cluster overload policy that never reports overload
+        ClusterOverloadPolicy mockPolicy = new ClusterOverloadPolicy()
+        {
+            @Override
+            public boolean isClusterOverloaded(InternalNodeManager nodeManager)
+            {
+                return false; // Never overloaded for tests
+            }
+
+            @Override
+            public String getName()
+            {
+                return "test-policy";
+            }
+        };
+
+        // Create a config with throttling disabled for tests
+        ClusterOverloadConfig config = new ClusterOverloadConfig()
+                .setClusterOverloadThrottlingEnabled(false);
+
+        return new ClusterResourceChecker(mockPolicy, config, new InMemoryNodeManager());
+    }
 }
diff --git a/presto-main-base/src/test/java/com/facebook/presto/sql/TestExpressionInterpreter.java b/presto-main-base/src/test/java/com/facebook/presto/sql/TestExpressionInterpreter.java
index f07caf36cf5e8..c4d9d050dbc3d 100644
--- a/presto-main-base/src/test/java/com/facebook/presto/sql/TestExpressionInterpreter.java
+++ b/presto-main-base/src/test/java/com/facebook/presto/sql/TestExpressionInterpreter.java
@@ -51,6 +51,7 @@
 import com.facebook.presto.sql.planner.Symbol;
 import com.facebook.presto.sql.planner.TypeProvider;
 import com.facebook.presto.sql.relational.FunctionResolution;
+import com.facebook.presto.sql.tree.EnumLiteral;
 import com.facebook.presto.sql.tree.Expression;
 import com.facebook.presto.sql.tree.ExpressionRewriter;
 import com.facebook.presto.sql.tree.ExpressionTreeRewriter;
@@ -899,6 +900,20 @@ public void testReservedWithDoubleQuotes()
         assertOptimizedEquals("\"time\"", "\"time\"");
     }
 
+    @Test
+    public void testEnumLiteralFormattingWithTypeAndValue()
+    {
+        java.util.function.BiFunction createEnumLiteral = (type, value) -> new EnumLiteral(Optional.empty(), type, value);
+        assertEquals(ExpressionFormatter.formatExpression(createEnumLiteral.apply("color", "RED"), Optional.empty()), "color: RED");
+        assertEquals(ExpressionFormatter.formatExpression(createEnumLiteral.apply("level", 1), Optional.empty()), "level: 1");
+        assertEquals(ExpressionFormatter.formatExpression(createEnumLiteral.apply("StatusType", "Active"), Optional.empty()), "StatusType: Active");
+        assertEquals(ExpressionFormatter.formatExpression(createEnumLiteral.apply("priority", "HIGH PRIORITY"), Optional.empty()), "priority: HIGH PRIORITY");
+        assertEquals(ExpressionFormatter.formatExpression(createEnumLiteral.apply("lang", "枚举"), Optional.empty()), "lang: 枚举");
+        assertEquals(ExpressionFormatter.formatExpression(createEnumLiteral.apply("special", "DOLLAR$"), Optional.empty()), "special: DOLLAR$");
+        assertEquals(ExpressionFormatter.formatExpression(createEnumLiteral.apply("enum_type", "VALUE_1"), Optional.empty()), "enum_type: VALUE_1");
+        assertEquals(ExpressionFormatter.formatExpression(createEnumLiteral.apply("flag", true), Optional.empty()), "flag: true");
+    }
+
     @Test
     public void testSearchCase()
     {
diff --git a/presto-main-base/src/test/java/com/facebook/presto/sql/analyzer/AbstractAnalyzerTest.java b/presto-main-base/src/test/java/com/facebook/presto/sql/analyzer/AbstractAnalyzerTest.java
index 3702dc249a00a..5819eadadcc6a 100644
--- a/presto-main-base/src/test/java/com/facebook/presto/sql/analyzer/AbstractAnalyzerTest.java
+++ b/presto-main-base/src/test/java/com/facebook/presto/sql/analyzer/AbstractAnalyzerTest.java
@@ -22,6 +22,10 @@
 import com.facebook.presto.common.type.StandardTypes;
 import com.facebook.presto.connector.informationSchema.InformationSchemaConnector;
 import com.facebook.presto.connector.system.SystemConnector;
+import com.facebook.presto.connector.tvf.TestingTableFunctions.DescriptorArgumentFunction;
+import com.facebook.presto.connector.tvf.TestingTableFunctions.SimpleTableFunction;
+import com.facebook.presto.connector.tvf.TestingTableFunctions.TableArgumentFunction;
+import com.facebook.presto.connector.tvf.TestingTableFunctions.TwoScalarArgumentsFunction;
 import com.facebook.presto.execution.warnings.WarningCollectorConfig;
 import com.facebook.presto.functionNamespace.SqlInvokedFunctionNamespaceManagerConfig;
 import com.facebook.presto.functionNamespace.execution.NoopSqlFunctionExecutor;
@@ -150,6 +154,13 @@ public void setup()
 
         metadata.getFunctionAndTypeManager().createFunction(SQL_FUNCTION_SQUARE, true);
 
+        metadata.getFunctionAndTypeManager().getTableFunctionRegistry().addTableFunctions(TPCH_CONNECTOR_ID,
+                ImmutableList.of(
+                        new SimpleTableFunction(),
+                        new TwoScalarArgumentsFunction(),
+                        new TableArgumentFunction(),
+                        new DescriptorArgumentFunction()));
+
         Catalog tpchTestCatalog = createTestingCatalog(TPCH_CATALOG, TPCH_CONNECTOR_ID);
         catalogManager.registerCatalog(tpchTestCatalog);
         metadata.getAnalyzePropertyManager().addProperties(TPCH_CONNECTOR_ID, tpchTestCatalog.getConnector(TPCH_CONNECTOR_ID).getAnalyzeProperties());
diff --git a/presto-main-base/src/test/java/com/facebook/presto/sql/analyzer/TestAnalyzer.java b/presto-main-base/src/test/java/com/facebook/presto/sql/analyzer/TestAnalyzer.java
index 48f442cf750ca..816bcd2a0aa5c 100644
--- a/presto-main-base/src/test/java/com/facebook/presto/sql/analyzer/TestAnalyzer.java
+++ b/presto-main-base/src/test/java/com/facebook/presto/sql/analyzer/TestAnalyzer.java
@@ -44,6 +44,10 @@
 import static com.facebook.presto.sql.analyzer.SemanticErrorCode.DUPLICATE_COLUMN_NAME;
 import static com.facebook.presto.sql.analyzer.SemanticErrorCode.DUPLICATE_PROPERTY;
 import static com.facebook.presto.sql.analyzer.SemanticErrorCode.DUPLICATE_RELATION;
+import static com.facebook.presto.sql.analyzer.SemanticErrorCode.EXPRESSION_NOT_CONSTANT;
+import static com.facebook.presto.sql.analyzer.SemanticErrorCode.FUNCTION_NOT_FOUND;
+import static com.facebook.presto.sql.analyzer.SemanticErrorCode.INVALID_ARGUMENTS;
+import static com.facebook.presto.sql.analyzer.SemanticErrorCode.INVALID_FUNCTION_ARGUMENT;
 import static com.facebook.presto.sql.analyzer.SemanticErrorCode.INVALID_FUNCTION_NAME;
 import static com.facebook.presto.sql.analyzer.SemanticErrorCode.INVALID_LITERAL;
 import static com.facebook.presto.sql.analyzer.SemanticErrorCode.INVALID_OFFSET_ROW_COUNT;
@@ -55,6 +59,7 @@
 import static com.facebook.presto.sql.analyzer.SemanticErrorCode.INVALID_WINDOW_FRAME;
 import static com.facebook.presto.sql.analyzer.SemanticErrorCode.MISMATCHED_COLUMN_ALIASES;
 import static com.facebook.presto.sql.analyzer.SemanticErrorCode.MISMATCHED_SET_COLUMN_TYPES;
+import static com.facebook.presto.sql.analyzer.SemanticErrorCode.MISSING_ARGUMENT;
 import static com.facebook.presto.sql.analyzer.SemanticErrorCode.MISSING_ATTRIBUTE;
 import static com.facebook.presto.sql.analyzer.SemanticErrorCode.MISSING_CATALOG;
 import static com.facebook.presto.sql.analyzer.SemanticErrorCode.MISSING_COLUMN;
@@ -1942,4 +1947,70 @@ public void testInvalidTemporaryFunctionName()
         assertFails(INVALID_FUNCTION_NAME, "CREATE TEMPORARY FUNCTION sum() RETURNS INT RETURN 1");
         assertFails(INVALID_FUNCTION_NAME, "CREATE TEMPORARY FUNCTION dev.test.foo() RETURNS INT RETURN 1");
     }
+
+    @Test
+    public void testTableFunctionNotFound()
+    {
+        assertFails(FUNCTION_NOT_FOUND,
+                "line 1:21: Table function non_existent_table_function not registered",
+                "SELECT * FROM TABLE(non_existent_table_function())");
+    }
+
+    @Test
+    public void testTableFunctionArguments()
+    {
+        assertFails(INVALID_ARGUMENTS, "line 1:51: Too many arguments. Expected at most 2 arguments, got 3 arguments", "SELECT * FROM TABLE(system.two_arguments_function(1, 2, 3))");
+
+        analyze("SELECT * FROM TABLE(system.two_arguments_function('foo'))");
+        analyze("SELECT * FROM TABLE(system.two_arguments_function(text => 'foo'))");
+        analyze("SELECT * FROM TABLE(system.two_arguments_function('foo', 1))");
+        analyze("SELECT * FROM TABLE(system.two_arguments_function(text => 'foo', number => 1))");
+
+        assertFails(INVALID_ARGUMENTS,
+                "line 1:51: All arguments must be passed by name or all must be passed positionally",
+                "SELECT * FROM TABLE(system.two_arguments_function('foo', number => 1))");
+
+        assertFails(INVALID_ARGUMENTS,
+                "line 1:51: All arguments must be passed by name or all must be passed positionally",
+                "SELECT * FROM TABLE(system.two_arguments_function(text => 'foo', 1))");
+
+        assertFails(INVALID_FUNCTION_ARGUMENT,
+                "line 1:66: Duplicate argument name: TEXT",
+                "SELECT * FROM TABLE(system.two_arguments_function(text => 'foo', text => 'bar'))");
+
+        // argument names are resolved in the canonical form
+        assertFails(INVALID_FUNCTION_ARGUMENT,
+                "line 1:66: Duplicate argument name: TEXT",
+                "SELECT * FROM TABLE(system.two_arguments_function(text => 'foo', TeXt => 'bar'))");
+
+        assertFails(INVALID_FUNCTION_ARGUMENT,
+                "line 1:66: Unexpected argument name: BAR",
+                "SELECT * FROM TABLE(system.two_arguments_function(text => 'foo', bar => 'bar'))");
+
+        assertFails(MISSING_ARGUMENT,
+                "line 1:51: Missing argument: TEXT",
+                "SELECT * FROM TABLE(system.two_arguments_function(number => 1))");
+    }
+
+    @Test
+    public void testScalarArgument()
+    {
+        analyze("SELECT * FROM TABLE(system.two_arguments_function('foo', 1))");
+
+        assertFails(INVALID_FUNCTION_ARGUMENT,
+                "line 1:64: Invalid argument NUMBER. Expected expression, got descriptor",
+                "SELECT * FROM TABLE(system.two_arguments_function(text => 'a', number => DESCRIPTOR(x integer, y boolean)))");
+
+        assertFails(INVALID_FUNCTION_ARGUMENT,
+                "line 1:64: 'descriptor' function is not allowed as a table function argument",
+                "SELECT * FROM TABLE(system.two_arguments_function(text => 'a', number => DESCRIPTOR(1 + 2)))");
+
+        assertFails(INVALID_FUNCTION_ARGUMENT,
+                "line 1:64: Invalid argument NUMBER. Expected expression, got table",
+                "SELECT * FROM TABLE(system.two_arguments_function(text => 'a', number => TABLE(t1)))");
+
+        assertFails(EXPRESSION_NOT_CONSTANT,
+                "line 1:74: Constant expression cannot contain a subquery",
+                "SELECT * FROM TABLE(system.two_arguments_function(text => 'a', number => (SELECT 1)))");
+    }
 }
diff --git a/presto-main-base/src/test/java/com/facebook/presto/sql/analyzer/TestFeaturesConfig.java b/presto-main-base/src/test/java/com/facebook/presto/sql/analyzer/TestFeaturesConfig.java
index 25a802b1e06ea..56be62a79306a 100644
--- a/presto-main-base/src/test/java/com/facebook/presto/sql/analyzer/TestFeaturesConfig.java
+++ b/presto-main-base/src/test/java/com/facebook/presto/sql/analyzer/TestFeaturesConfig.java
@@ -201,6 +201,7 @@ public void testDefaults()
                 .setPushRemoteExchangeThroughGroupId(false)
                 .setOptimizeMultipleApproxPercentileOnSameFieldEnabled(true)
                 .setNativeExecutionEnabled(false)
+                .setBuiltInSidecarFunctionsEnabled(false)
                 .setDisableTimeStampWithTimeZoneForNative(false)
                 .setDisableIPAddressForNative(false)
                 .setNativeExecutionExecutablePath("./presto_server")
@@ -262,6 +263,7 @@ public void testDefaults()
                 .setAddExchangeBelowPartialAggregationOverGroupId(false)
                 .setAddDistinctBelowSemiJoinBuild(false)
                 .setPushdownSubfieldForMapFunctions(true)
+                .setUtilizeUniquePropertyInQueryPlanning(true)
                 .setInnerJoinPushdownEnabled(false)
                 .setBroadcastSemiJoinForDelete(true)
                 .setInEqualityJoinPushdownEnabled(false)
@@ -416,6 +418,7 @@ public void testExplicitPropertyMappings()
                 .put("optimizer.push-remote-exchange-through-group-id", "true")
                 .put("optimizer.optimize-multiple-approx-percentile-on-same-field", "false")
                 .put("native-execution-enabled", "true")
+                .put("built-in-sidecar-functions-enabled", "true")
                 .put("disable-timestamp-with-timezone-for-native-execution", "true")
                 .put("disable-ipaddress-for-native-execution", "true")
                 .put("native-execution-executable-path", "/bin/echo")
@@ -480,6 +483,7 @@ public void testExplicitPropertyMappings()
                 .put("exclude-invalid-worker-session-properties", "true")
                 .put("optimizer.add-distinct-below-semi-join-build", "true")
                 .put("optimizer.pushdown-subfield-for-map-functions", "false")
+                .put("optimizer.utilize-unique-property-in-query-planning", "false")
                 .put("optimizer.add-exchange-below-partial-aggregation-over-group-id", "true")
                 .put("max_serializable_object_size", "50")
                 .build();
@@ -628,6 +632,7 @@ public void testExplicitPropertyMappings()
                 .setPushRemoteExchangeThroughGroupId(true)
                 .setOptimizeMultipleApproxPercentileOnSameFieldEnabled(false)
                 .setNativeExecutionEnabled(true)
+                .setBuiltInSidecarFunctionsEnabled(true)
                 .setDisableTimeStampWithTimeZoneForNative(true)
                 .setDisableIPAddressForNative(true)
                 .setNativeExecutionExecutablePath("/bin/echo")
@@ -689,6 +694,7 @@ public void testExplicitPropertyMappings()
                 .setAddExchangeBelowPartialAggregationOverGroupId(true)
                 .setAddDistinctBelowSemiJoinBuild(true)
                 .setPushdownSubfieldForMapFunctions(false)
+                .setUtilizeUniquePropertyInQueryPlanning(false)
                 .setInEqualityJoinPushdownEnabled(true)
                 .setBroadcastSemiJoinForDelete(false)
                 .setRewriteMinMaxByToTopNEnabled(true)
diff --git a/presto-main-base/src/test/java/com/facebook/presto/sql/planner/TestingWriterTarget.java b/presto-main-base/src/test/java/com/facebook/presto/sql/planner/TestingWriterTarget.java
index 863047690952b..efb9cfdccd6ec 100644
--- a/presto-main-base/src/test/java/com/facebook/presto/sql/planner/TestingWriterTarget.java
+++ b/presto-main-base/src/test/java/com/facebook/presto/sql/planner/TestingWriterTarget.java
@@ -14,9 +14,17 @@
 
 package com.facebook.presto.sql.planner;
 
+import com.facebook.presto.common.QualifiedObjectName;
+import com.facebook.presto.common.SourceColumn;
 import com.facebook.presto.spi.ConnectorId;
 import com.facebook.presto.spi.SchemaTableName;
+import com.facebook.presto.spi.eventlistener.OutputColumnMetadata;
 import com.facebook.presto.spi.plan.TableWriterNode;
+import com.google.common.collect.ImmutableList;
+import com.google.common.collect.ImmutableSet;
+
+import java.util.List;
+import java.util.Optional;
 
 public class TestingWriterTarget
         extends TableWriterNode.WriterTarget
@@ -36,6 +44,17 @@ public SchemaTableName getSchemaTableName()
         return SCHEMA_TABLE_NAME;
     }
 
+    @Override
+    public Optional> getOutputColumns()
+    {
+        return Optional.of(
+                ImmutableList.of(
+                        new OutputColumnMetadata(
+                                "column", "type",
+                                ImmutableSet.of(
+                                        new SourceColumn(QualifiedObjectName.valueOf("catalog.schema.table"), "column")))));
+    }
+
     @Override
     public String toString()
     {
diff --git a/presto-main-base/src/test/java/com/facebook/presto/sql/planner/assertions/PlanMatchPattern.java b/presto-main-base/src/test/java/com/facebook/presto/sql/planner/assertions/PlanMatchPattern.java
index 9c6eb9a528e56..c69c4ea815c92 100644
--- a/presto-main-base/src/test/java/com/facebook/presto/sql/planner/assertions/PlanMatchPattern.java
+++ b/presto-main-base/src/test/java/com/facebook/presto/sql/planner/assertions/PlanMatchPattern.java
@@ -412,6 +412,11 @@ public static PlanMatchPattern strictProject(Map assi
                 .withExactAssignments(assignments.values());
     }
 
+    public static PlanMatchPattern semiJoin(PlanMatchPattern source, PlanMatchPattern filtering)
+    {
+        return node(SemiJoinNode.class, source, filtering);
+    }
+
     public static PlanMatchPattern semiJoin(String sourceSymbolAlias, String filteringSymbolAlias, String outputAlias, PlanMatchPattern source, PlanMatchPattern filtering)
     {
         return semiJoin(sourceSymbolAlias, filteringSymbolAlias, outputAlias, Optional.empty(), source, filtering);
diff --git a/presto-main-base/src/test/java/com/facebook/presto/sql/planner/iterative/rule/TestCrossJoinWithArrayNotContainsToAntiJoin.java b/presto-main-base/src/test/java/com/facebook/presto/sql/planner/iterative/rule/TestCrossJoinWithArrayNotContainsToAntiJoin.java
index 4bd3f167ac11a..194cf489d3998 100644
--- a/presto-main-base/src/test/java/com/facebook/presto/sql/planner/iterative/rule/TestCrossJoinWithArrayNotContainsToAntiJoin.java
+++ b/presto-main-base/src/test/java/com/facebook/presto/sql/planner/iterative/rule/TestCrossJoinWithArrayNotContainsToAntiJoin.java
@@ -14,6 +14,7 @@
 package com.facebook.presto.sql.planner.iterative.rule;
 
 import com.facebook.presto.common.type.ArrayType;
+import com.facebook.presto.scalar.sql.SqlInvokedFunctionsPlugin;
 import com.facebook.presto.spi.plan.Assignments;
 import com.facebook.presto.spi.plan.FilterNode;
 import com.facebook.presto.spi.plan.JoinNode;
@@ -22,9 +23,11 @@
 import com.facebook.presto.spi.plan.UnnestNode;
 import com.facebook.presto.spi.plan.ValuesNode;
 import com.facebook.presto.sql.planner.iterative.rule.test.BaseRuleTest;
+import com.facebook.presto.sql.planner.iterative.rule.test.RuleTester;
 import com.facebook.presto.sql.planner.plan.EnforceSingleRowNode;
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableSet;
+import org.testng.annotations.BeforeClass;
 import org.testng.annotations.Test;
 
 import static com.facebook.presto.SystemSessionProperties.REWRITE_CROSS_JOIN_ARRAY_NOT_CONTAINS_TO_ANTI_JOIN;
@@ -33,10 +36,18 @@
 import static com.facebook.presto.common.type.VarcharType.VARCHAR;
 import static com.facebook.presto.sql.planner.assertions.PlanMatchPattern.node;
 import static com.facebook.presto.sql.planner.iterative.rule.test.PlanBuilder.constantExpressions;
+import static java.util.Collections.singletonList;
 
 public class TestCrossJoinWithArrayNotContainsToAntiJoin
         extends BaseRuleTest
 {
+    @BeforeClass
+    @Override
+    public void setUp()
+    {
+        tester = new RuleTester(singletonList(new SqlInvokedFunctionsPlugin()));
+    }
+
     @Test
     public void testTriggerForBigInt()
     {
diff --git a/presto-main-base/src/test/java/com/facebook/presto/sql/planner/iterative/rule/TestLeftJoinWithArrayContainsToEquiJoinCondition.java b/presto-main-base/src/test/java/com/facebook/presto/sql/planner/iterative/rule/TestLeftJoinWithArrayContainsToEquiJoinCondition.java
index 657093aec5be2..cdf060f86d390 100644
--- a/presto-main-base/src/test/java/com/facebook/presto/sql/planner/iterative/rule/TestLeftJoinWithArrayContainsToEquiJoinCondition.java
+++ b/presto-main-base/src/test/java/com/facebook/presto/sql/planner/iterative/rule/TestLeftJoinWithArrayContainsToEquiJoinCondition.java
@@ -14,11 +14,14 @@
 package com.facebook.presto.sql.planner.iterative.rule;
 
 import com.facebook.presto.common.type.ArrayType;
+import com.facebook.presto.scalar.sql.SqlInvokedFunctionsPlugin;
 import com.facebook.presto.spi.plan.JoinType;
 import com.facebook.presto.sql.planner.iterative.rule.test.BaseRuleTest;
+import com.facebook.presto.sql.planner.iterative.rule.test.RuleTester;
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableMap;
 import com.google.common.collect.ImmutableSet;
+import org.testng.annotations.BeforeClass;
 import org.testng.annotations.Test;
 
 import java.util.Optional;
@@ -32,10 +35,18 @@
 import static com.facebook.presto.sql.planner.assertions.PlanMatchPattern.project;
 import static com.facebook.presto.sql.planner.assertions.PlanMatchPattern.unnest;
 import static com.facebook.presto.sql.planner.assertions.PlanMatchPattern.values;
+import static java.util.Collections.singletonList;
 
 public class TestLeftJoinWithArrayContainsToEquiJoinCondition
         extends BaseRuleTest
 {
+    @BeforeClass
+    @Override
+    public void setUp()
+    {
+        tester = new RuleTester(singletonList(new SqlInvokedFunctionsPlugin()));
+    }
+
     @Test
     public void testTriggerForBigIntArrayRightSide()
     {
diff --git a/presto-main-base/src/test/java/com/facebook/presto/sql/planner/optimizations/TestLocalProperties.java b/presto-main-base/src/test/java/com/facebook/presto/sql/planner/optimizations/TestLocalProperties.java
index db0600d2d6380..edbeeb3e0cdc2 100644
--- a/presto-main-base/src/test/java/com/facebook/presto/sql/planner/optimizations/TestLocalProperties.java
+++ b/presto-main-base/src/test/java/com/facebook/presto/sql/planner/optimizations/TestLocalProperties.java
@@ -20,6 +20,7 @@
 import com.facebook.presto.spi.GroupingProperty;
 import com.facebook.presto.spi.LocalProperty;
 import com.facebook.presto.spi.SortingProperty;
+import com.facebook.presto.spi.UniqueProperty;
 import com.facebook.presto.spi.relation.VariableReferenceExpression;
 import com.facebook.presto.testing.TestingMetadata.TestingColumnHandle;
 import com.fasterxml.jackson.core.JsonParser;
@@ -78,6 +79,10 @@ public void testConstantProcessing()
         input = ImmutableList.of(constant("a"), constant("b"));
         assertEquals(stripLeadingConstants(input), ImmutableList.of());
         assertEquals(extractLeadingConstants(input), ImmutableSet.of("a", "b"));
+
+        input = ImmutableList.of(unique("a"));
+        assertEquals(stripLeadingConstants(input), ImmutableList.of(unique("a")));
+        assertEquals(extractLeadingConstants(input), ImmutableSet.of());
     }
 
     @Test
@@ -138,6 +143,14 @@ public void testTranslate()
         map = ImmutableMap.of("a", "a1", "b", "b1", "c", "c1");
         input = ImmutableList.of(grouped("a"), constant("b"), grouped("c"));
         assertEquals(LocalProperties.translate(input, translateWithMap(map)), ImmutableList.of(grouped("a1"), constant("b1"), grouped("c1")));
+
+        map = ImmutableMap.of();
+        input = ImmutableList.of(unique("a"));
+        assertEquals(LocalProperties.translate(input, translateWithMap(map)), ImmutableList.of());
+
+        map = ImmutableMap.of("a", "a1");
+        input = ImmutableList.of(unique("a"));
+        assertEquals(LocalProperties.translate(input, translateWithMap(map)), ImmutableList.of(unique("a1")));
     }
 
     private static  Function> translateWithMap(Map translateMap)
@@ -177,6 +190,35 @@ public void testNormalizeOverlappingSymbol()
         assertNormalizeAndFlatten(
                 localProperties,
                 grouped("a"));
+
+        localProperties = builder()
+                .unique("a")
+                .sorted("a", SortOrder.ASC_NULLS_FIRST)
+                .constant("a")
+                .build();
+        assertNormalize(
+                localProperties,
+                Optional.of(unique("a")),
+                Optional.empty(),
+                Optional.empty());
+        assertNormalizeAndFlatten(
+                localProperties,
+                unique("a"));
+
+        localProperties = builder()
+                .grouped("a")
+                .unique("a")
+                .constant("a")
+                .build();
+        assertNormalize(
+                localProperties,
+                Optional.of(grouped("a")),
+                Optional.of(unique("a")),
+                Optional.empty());
+        assertNormalizeAndFlatten(
+                localProperties,
+                grouped("a"),
+                unique("a"));
     }
 
     @Test
@@ -780,6 +822,11 @@ private static GroupingProperty grouped(String... columns)
         return new GroupingProperty<>(Arrays.asList(columns));
     }
 
+    private static UniqueProperty unique(String column)
+    {
+        return new UniqueProperty<>(column);
+    }
+
     private static SortingProperty sorted(String column, SortOrder order)
     {
         return new SortingProperty<>(column, order);
@@ -812,6 +859,12 @@ public Builder constant(String column)
             return this;
         }
 
+        public Builder unique(String column)
+        {
+            properties.add(new UniqueProperty<>(column));
+            return this;
+        }
+
         public List> build()
         {
             return new ArrayList<>(properties);
diff --git a/presto-main-base/src/test/java/com/facebook/presto/tdigest/TestTDigest.java b/presto-main-base/src/test/java/com/facebook/presto/tdigest/TestTDigest.java
index 773be051a5df5..04320aaa52fa8 100644
--- a/presto-main-base/src/test/java/com/facebook/presto/tdigest/TestTDigest.java
+++ b/presto-main-base/src/test/java/com/facebook/presto/tdigest/TestTDigest.java
@@ -397,6 +397,16 @@ public void testGeometricDistribution()
         }
     }
 
+    @Test(expectedExceptions = IllegalArgumentException.class, expectedExceptionsMessageRegExp = ".*Cannot serialize t-digest with NaN mean value.*")
+    public void testSerializationWithNaNMean()
+    {
+        double[] means = {1.0, Double.NaN, 3.0};
+        double[] weights = {1.0, 1.0, 1.0};
+        TDigest tDigest = createTDigest(means, weights, STANDARD_COMPRESSION_FACTOR, 1.0, 3.0, 6.0, 3);
+
+        tDigest.serialize();
+    }
+
     @Test(enabled = false)
     public void testPoissonDistribution()
     {
diff --git a/presto-main-base/src/test/java/com/facebook/presto/type/TestBuiltInTypeRegistry.java b/presto-main-base/src/test/java/com/facebook/presto/type/TestBuiltInTypeRegistry.java
index b5a4b8f49f3d5..0c626e5f0fc3d 100644
--- a/presto-main-base/src/test/java/com/facebook/presto/type/TestBuiltInTypeRegistry.java
+++ b/presto-main-base/src/test/java/com/facebook/presto/type/TestBuiltInTypeRegistry.java
@@ -13,11 +13,11 @@
  */
 package com.facebook.presto.type;
 
-import com.facebook.presto.UnknownTypeException;
 import com.facebook.presto.common.type.Type;
 import com.facebook.presto.common.type.TypeSignature;
 import com.facebook.presto.metadata.FunctionAndTypeManager;
 import com.facebook.presto.metadata.OperatorNotFoundException;
+import com.facebook.presto.spi.type.UnknownTypeException;
 import com.google.common.collect.ImmutableSet;
 import org.testng.annotations.Test;
 
diff --git a/presto-main-base/src/test/java/com/facebook/presto/type/TestMapOperators.java b/presto-main-base/src/test/java/com/facebook/presto/type/TestMapOperators.java
index 0690252130ce8..c00ad1b5baf86 100644
--- a/presto-main-base/src/test/java/com/facebook/presto/type/TestMapOperators.java
+++ b/presto-main-base/src/test/java/com/facebook/presto/type/TestMapOperators.java
@@ -527,16 +527,13 @@ public void testJsonToMap()
                         .put("k8", "[null]")
                         .build());
 
-        // These two tests verifies that partial json cast preserves input order
-        // The second test should never happen in real life because valid json in presto requires natural key ordering.
-        // However, it is added to make sure that the order in the first test is not a coincidence.
         assertFunction("CAST(JSON '{\"k1\": {\"1klmnopq\":1, \"2klmnopq\":2, \"3klmnopq\":3, \"4klmnopq\":4, \"5klmnopq\":5, \"6klmnopq\":6, \"7klmnopq\":7}}' AS MAP)",
                 mapType(VARCHAR, JSON),
                 ImmutableMap.of("k1", "{\"1klmnopq\":1,\"2klmnopq\":2,\"3klmnopq\":3,\"4klmnopq\":4,\"5klmnopq\":5,\"6klmnopq\":6,\"7klmnopq\":7}"));
+
         assertFunction("CAST(unchecked_to_json('{\"k1\": {\"7klmnopq\":7, \"6klmnopq\":6, \"5klmnopq\":5, \"4klmnopq\":4, \"3klmnopq\":3, \"2klmnopq\":2, \"1klmnopq\":1}}') AS MAP)",
                 mapType(VARCHAR, JSON),
-                ImmutableMap.of("k1", "{\"7klmnopq\":7,\"6klmnopq\":6,\"5klmnopq\":5,\"4klmnopq\":4,\"3klmnopq\":3,\"2klmnopq\":2,\"1klmnopq\":1}"));
-
+                ImmutableMap.of("k1", "{\"1klmnopq\":1,\"2klmnopq\":2,\"3klmnopq\":3,\"4klmnopq\":4,\"5klmnopq\":5,\"6klmnopq\":6,\"7klmnopq\":7}"));
         // nested array/map
         assertFunction("CAST(JSON '{\"1\": [1, 2], \"2\": [3, null], \"3\": [], \"5\": [null, null], \"8\": null}' AS MAP>)",
                 mapType(BIGINT, new ArrayType(BIGINT)),
diff --git a/presto-main/etc/config.properties b/presto-main/etc/config.properties
index a120422c4b6f1..164015db096d4 100644
--- a/presto-main/etc/config.properties
+++ b/presto-main/etc/config.properties
@@ -51,7 +51,8 @@ plugin.bundles=\
   ../presto-node-ttl-fetchers/pom.xml,\
   ../presto-hive-function-namespace/pom.xml,\
   ../presto-delta/pom.xml,\
-  ../presto-hudi/pom.xml
+  ../presto-hudi/pom.xml, \
+  ../presto-sql-invoked-functions-plugin/pom.xml
 
 presto.version=testversion
 node-scheduler.include-coordinator=true
diff --git a/presto-main/etc/regex-map.txt b/presto-main/etc/regex-map.txt
new file mode 100644
index 0000000000000..b61daa86bcc42
--- /dev/null
+++ b/presto-main/etc/regex-map.txt
@@ -0,0 +1,3 @@
+user=.*
+internal=coordinator
+admin=su.*
\ No newline at end of file
diff --git a/presto-main/pom.xml b/presto-main/pom.xml
index 07e0f4828ceb0..5f737a7ab3bd6 100644
--- a/presto-main/pom.xml
+++ b/presto-main/pom.xml
@@ -23,6 +23,11 @@
             presto-main-base
         
 
+        
+            com.facebook.presto
+            presto-built-in-worker-function-tools
+        
+
         
             com.facebook.airlift
             jmx-http
@@ -53,6 +58,11 @@
             presto-common
         
 
+        
+            com.facebook.presto
+            presto-function-namespace-managers-common
+        
+
         
             io.jsonwebtoken
             jjwt-api
@@ -147,98 +157,122 @@
             com.facebook.airlift.drift
             drift-codec
         
+
         
             com.google.guava
             guava
         
+
         
             com.google.errorprone
             error_prone_annotations
         
+
         
             com.facebook.airlift.drift
             drift-client
         
+
         
             com.facebook.airlift.drift
             drift-transport-netty
         
+
         
             io.netty
             netty-buffer
         
+
         
             com.facebook.airlift.drift
             drift-api
         
+
         
             com.facebook.airlift
             http-client
         
+
         
             joda-time
             joda-time
         
+
         
             com.facebook.airlift
             concurrent
         
+
         
             io.airlift
             slice
         
+
         
             com.facebook.presto
             presto-analyzer
         
+
         
             com.facebook.airlift
             units
         
+
         
             com.facebook.presto
             presto-client
         
+
         
             com.facebook.airlift.drift
             drift-protocol
         
+
         
             com.facebook.airlift
             stats
         
+
         
             com.facebook.presto
             presto-memory-context
         
+
         
             jakarta.inject
             jakarta.inject-api
         
+
         
             jakarta.servlet
             jakarta.servlet-api
         
+
         
             com.facebook.presto
             presto-spi
         
+
         
             com.facebook.airlift
             trace-token
         
+
         
             com.facebook.airlift
             http-server
         
+
         
             com.facebook.airlift
             event
         
+
         
             com.facebook.airlift
             bootstrap
         
+
         
             com.facebook.presto
             presto-parser
@@ -273,7 +307,7 @@
             io.projectreactor.netty
             reactor-netty-http
         
-        
+
         
             io.micrometer
             micrometer-core
@@ -339,6 +373,47 @@
             runtime
         
 
+        
+            org.apache.commons
+            commons-lang3
+        
+
+        
+            com.nimbusds
+            nimbus-jose-jwt
+        
+
+        
+            com.nimbusds
+            oauth2-oidc-sdk
+            
+                
+                    org.ow2.asm
+                    asm
+                
+            
+        
+
+        
+            com.fasterxml.jackson.core
+            jackson-core
+        
+
+        
+            io.airlift
+            aircompressor
+        
+
+        
+            net.jodah
+            failsafe
+        
+
+        
+            net.minidev
+            json-smart
+        
+
         
         
             com.facebook.presto
@@ -405,6 +480,18 @@
                 
             
         
+
+        
+            org.testcontainers
+            testcontainers
+            test
+        
+
+        
+            org.testcontainers
+            postgresql
+            test
+        
     
 
     
@@ -440,7 +527,13 @@
                         com.facebook.airlift.drift:drift-protocol
                         io.netty:netty-buffer
                         javax.inject:javax.inject
+                        com.squareup.okhttp3:okhttp
+                        com.squareup.okhttp3:okhttp-urlconnection
                     
+                    
+                        com.squareup.okhttp3:okhttp
+                        com.squareup.okhttp3:okhttp-urlconnection
+                    
                 
             
         
diff --git a/presto-main/src/main/java/com/facebook/presto/metadata/DiscoveryNodeManager.java b/presto-main/src/main/java/com/facebook/presto/metadata/DiscoveryNodeManager.java
index 02fdf0efe62cf..99d9a8d92dfd3 100644
--- a/presto-main/src/main/java/com/facebook/presto/metadata/DiscoveryNodeManager.java
+++ b/presto-main/src/main/java/com/facebook/presto/metadata/DiscoveryNodeManager.java
@@ -27,8 +27,10 @@
 import com.facebook.presto.server.InternalCommunicationConfig.CommunicationProtocol;
 import com.facebook.presto.server.thrift.ThriftServerInfoClient;
 import com.facebook.presto.spi.ConnectorId;
+import com.facebook.presto.spi.NodeLoadMetrics;
 import com.facebook.presto.spi.NodePoolType;
 import com.facebook.presto.spi.NodeState;
+import com.facebook.presto.spi.NodeStats;
 import com.facebook.presto.statusservice.NodeStatusService;
 import com.google.common.base.Splitter;
 import com.google.common.collect.HashMultimap;
@@ -93,7 +95,7 @@ public final class DiscoveryNodeManager
     private final FailureDetector failureDetector;
     private final Optional nodeStatusService;
     private final NodeVersion expectedNodeVersion;
-    private final ConcurrentHashMap nodeStates = new ConcurrentHashMap<>();
+    private final ConcurrentHashMap nodeStats = new ConcurrentHashMap<>();
     private final HttpClient httpClient;
     private final DriftClient driftClient;
     private final ScheduledExecutorService nodeStateUpdateExecutor;
@@ -102,6 +104,7 @@ public final class DiscoveryNodeManager
     private final InternalNode currentNode;
     private final CommunicationProtocol protocol;
     private final boolean isMemoizeDeadNodesEnabled;
+    private final InternalCommunicationConfig internalCommunicationConfig;
 
     @GuardedBy("this")
     private SetMultimap activeNodesByConnectorId;
@@ -153,6 +156,7 @@ public DiscoveryNodeManager(
         this.nodeStateUpdateExecutor = newSingleThreadScheduledExecutor(threadsNamed("node-state-poller-%s"));
         this.nodeStateEventExecutor = newCachedThreadPool(threadsNamed("node-state-events-%s"));
         this.httpsRequired = internalCommunicationConfig.isHttpsRequired();
+        this.internalCommunicationConfig = requireNonNull(internalCommunicationConfig, "internalCommunicationConfig is null");
 
         this.currentNode = findCurrentNode(
                 serviceSelector.selectAllServices(),
@@ -211,6 +215,7 @@ private static NodePoolType getPoolType(ServiceDescriptor service)
     @PostConstruct
     public void startPollingNodeStates()
     {
+        long pollingIntervalMillis = internalCommunicationConfig.getNodeDiscoveryPollingIntervalMillis();
         nodeStateUpdateExecutor.scheduleWithFixedDelay(() -> {
             try {
                 pollWorkers();
@@ -218,7 +223,7 @@ public void startPollingNodeStates()
             catch (Exception e) {
                 log.error(e, "Error polling state of nodes");
             }
-        }, 5, 5, TimeUnit.SECONDS);
+        }, pollingIntervalMillis, pollingIntervalMillis, TimeUnit.MILLISECONDS);
         pollWorkers();
     }
 
@@ -236,20 +241,20 @@ private void pollWorkers()
 
         // Remove nodes that don't exist anymore
         // Make a copy to materialize the set difference
-        Set deadNodes = difference(nodeStates.keySet(), aliveNodeIds).immutableCopy();
-        nodeStates.keySet().removeAll(deadNodes);
+        Set deadNodes = difference(nodeStats.keySet(), aliveNodeIds).immutableCopy();
+        nodeStats.keySet().removeAll(deadNodes);
 
         // Add new nodes
         for (InternalNode node : aliveNodes) {
             switch (protocol) {
                 case HTTP:
-                    nodeStates.putIfAbsent(node.getNodeIdentifier(),
-                            new HttpRemoteNodeState(httpClient, uriBuilderFrom(node.getInternalUri()).appendPath("/v1/info/state").build()));
+                    nodeStats.putIfAbsent(node.getNodeIdentifier(),
+                            new HttpRemoteNodeStats(httpClient, uriBuilderFrom(node.getInternalUri()).appendPath("/v1/info/stats").build(), internalCommunicationConfig.getNodeStatsRefreshIntervalMillis()));
                     break;
                 case THRIFT:
                     if (node.getThriftPort().isPresent()) {
-                        nodeStates.put(node.getNodeIdentifier(),
-                                new ThriftRemoteNodeState(driftClient, uriBuilderFrom(node.getInternalUri()).scheme("thrift").port(node.getThriftPort().getAsInt()).build()));
+                        nodeStats.put(node.getNodeIdentifier(),
+                                new ThriftRemoteNodeStats(driftClient, uriBuilderFrom(node.getInternalUri()).scheme("thrift").port(node.getThriftPort().getAsInt()).build(), internalCommunicationConfig.getNodeStatsRefreshIntervalMillis()));
                     }
                     else {
                         // thrift port has not yet been populated; ignore the node for now
@@ -259,7 +264,7 @@ private void pollWorkers()
         }
 
         // Schedule refresh
-        nodeStates.values().forEach(RemoteNodeState::asyncRefresh);
+        nodeStats.values().forEach(RemoteNodeStats::asyncRefresh);
 
         // update indexes
         refreshNodesInternal();
@@ -438,10 +443,19 @@ private NodeState getNodeState(InternalNode node)
 
     private boolean isNodeShuttingDown(String nodeId)
     {
-        Optional remoteNodeState = nodeStates.containsKey(nodeId)
-                ? nodeStates.get(nodeId).getNodeState()
+        Optional remoteNodeStats = nodeStats.containsKey(nodeId)
+                ? nodeStats.get(nodeId).getNodeStats()
                 : Optional.empty();
-        return remoteNodeState.isPresent() && remoteNodeState.get() == SHUTTING_DOWN;
+        return remoteNodeStats.isPresent() && remoteNodeStats.get().getNodeState() == SHUTTING_DOWN;
+    }
+
+    @Override
+    public Optional getNodeLoadMetrics(String nodeId)
+    {
+        Optional remoteNodeStats = nodeStats.containsKey(nodeId)
+                ? nodeStats.get(nodeId).getNodeStats()
+                : Optional.empty();
+        return remoteNodeStats.flatMap(NodeStats::getLoadMetrics);
     }
 
     @Override
diff --git a/presto-main/src/main/java/com/facebook/presto/metadata/HttpRemoteNodeStats.java b/presto-main/src/main/java/com/facebook/presto/metadata/HttpRemoteNodeStats.java
new file mode 100644
index 0000000000000..5f99a378c4d64
--- /dev/null
+++ b/presto-main/src/main/java/com/facebook/presto/metadata/HttpRemoteNodeStats.java
@@ -0,0 +1,121 @@
+/*
+ * Licensed 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 com.facebook.presto.metadata;
+
+import com.facebook.airlift.http.client.FullJsonResponseHandler.JsonResponse;
+import com.facebook.airlift.http.client.HttpClient;
+import com.facebook.airlift.http.client.HttpClient.HttpResponseFuture;
+import com.facebook.airlift.http.client.Request;
+import com.facebook.airlift.log.Logger;
+import com.facebook.airlift.units.Duration;
+import com.facebook.presto.spi.NodeStats;
+import com.google.common.util.concurrent.FutureCallback;
+import com.google.common.util.concurrent.Futures;
+import com.google.errorprone.annotations.ThreadSafe;
+import jakarta.annotation.Nullable;
+
+import java.net.URI;
+import java.util.Optional;
+import java.util.concurrent.Future;
+import java.util.concurrent.atomic.AtomicLong;
+import java.util.concurrent.atomic.AtomicReference;
+
+import static com.facebook.airlift.http.client.FullJsonResponseHandler.createFullJsonResponseHandler;
+import static com.facebook.airlift.http.client.HttpStatus.OK;
+import static com.facebook.airlift.http.client.Request.Builder.prepareGet;
+import static com.facebook.airlift.json.JsonCodec.jsonCodec;
+import static com.facebook.airlift.units.Duration.nanosSince;
+import static com.google.common.net.MediaType.JSON_UTF_8;
+import static com.google.common.util.concurrent.MoreExecutors.directExecutor;
+import static jakarta.ws.rs.core.HttpHeaders.CONTENT_TYPE;
+import static java.util.Objects.requireNonNull;
+import static java.util.concurrent.TimeUnit.SECONDS;
+
+@ThreadSafe
+public class HttpRemoteNodeStats
+        implements RemoteNodeStats
+{
+    private static final Logger log = Logger.get(HttpRemoteNodeStats.class);
+
+    private final HttpClient httpClient;
+    private final URI stateInfoUri;
+    private final long refreshIntervalMillis;
+    private final AtomicReference> nodeStats = new AtomicReference<>(Optional.empty());
+    private final AtomicReference> future = new AtomicReference<>();
+    private final AtomicLong lastUpdateNanos = new AtomicLong();
+    private final AtomicLong lastWarningLogged = new AtomicLong();
+
+    public HttpRemoteNodeStats(HttpClient httpClient, URI stateInfoUri, long refreshIntervalMillis)
+    {
+        this.httpClient = requireNonNull(httpClient, "httpClient is null");
+        this.stateInfoUri = requireNonNull(stateInfoUri, "stateInfoUri is null");
+        this.refreshIntervalMillis = refreshIntervalMillis;
+    }
+
+    @Override
+    public Optional getNodeStats()
+    {
+        return nodeStats.get();
+    }
+
+    @Override
+    public synchronized void asyncRefresh()
+    {
+        Duration sinceUpdate = nanosSince(lastUpdateNanos.get());
+        if (nanosSince(lastWarningLogged.get()).toMillis() > 1_000 &&
+                sinceUpdate.toMillis() > 10_000 &&
+                future.get() != null) {
+            log.warn("Node state update request to %s has not returned in %s", stateInfoUri, sinceUpdate.toString(SECONDS));
+            lastWarningLogged.set(System.nanoTime());
+        }
+        if (sinceUpdate.toMillis() > refreshIntervalMillis && future.get() == null) {
+            Request request = prepareGet()
+                    .setUri(stateInfoUri)
+                    .setHeader(CONTENT_TYPE, JSON_UTF_8.toString())
+                    .build();
+            HttpResponseFuture> responseFuture = httpClient.executeAsync(request, createFullJsonResponseHandler(jsonCodec(NodeStats.class)));
+            future.compareAndSet(null, responseFuture);
+
+            Futures.addCallback(responseFuture, new FutureCallback>()
+            {
+                @Override
+                public void onSuccess(@Nullable JsonResponse result)
+                {
+                    lastUpdateNanos.set(System.nanoTime());
+                    future.compareAndSet(responseFuture, null);
+                    if (result != null) {
+                        if (result.hasValue()) {
+                            nodeStats.set(Optional.ofNullable(result.getValue()));
+                        }
+                        if (result.getStatusCode() != OK.code()) {
+                            log.warn("Error fetching node stats from %s returned status %d", stateInfoUri, result.getStatusCode());
+                            return;
+                        }
+                    }
+                    else {
+                        log.warn("Node statistics endpoint %s returned null response, using cached statistics", stateInfoUri);
+                    }
+                }
+
+                @Override
+                public void onFailure(Throwable t)
+                {
+                    log.error("Error fetching node stats from %s: %s", stateInfoUri, t.getMessage());
+                    lastUpdateNanos.set(System.nanoTime());
+                    future.compareAndSet(responseFuture, null);
+                }
+            }, directExecutor());
+        }
+    }
+}
diff --git a/presto-main/src/main/java/com/facebook/presto/server/CoordinatorModule.java b/presto-main/src/main/java/com/facebook/presto/server/CoordinatorModule.java
index a165616c09f4c..1813c5eb06995 100644
--- a/presto-main/src/main/java/com/facebook/presto/server/CoordinatorModule.java
+++ b/presto-main/src/main/java/com/facebook/presto/server/CoordinatorModule.java
@@ -144,7 +144,7 @@ public class CoordinatorModule
 {
     private static final String DEFAULT_WEBUI_CSP =
             "default-src 'self'; style-src 'self' 'unsafe-inline' https://fonts.googleapis.com; " +
-                    "font-src 'self' https://fonts.gstatic.com; frame-ancestors 'self'; img-src http: https: data:";
+                    "font-src 'self' https://fonts.gstatic.com; frame-ancestors 'self'; img-src 'self' http: https: data:; form-action 'self'";
 
     public static HttpResourceBinding webUIBinder(Binder binder, String path, String classPathResourceBase)
     {
diff --git a/presto-main/src/main/java/com/facebook/presto/server/PrestoServer.java b/presto-main/src/main/java/com/facebook/presto/server/PrestoServer.java
index 552b6da5d812e..5c9f7bbe6099d 100644
--- a/presto-main/src/main/java/com/facebook/presto/server/PrestoServer.java
+++ b/presto-main/src/main/java/com/facebook/presto/server/PrestoServer.java
@@ -33,6 +33,7 @@
 import com.facebook.drift.transport.netty.server.DriftNettyServerTransport;
 import com.facebook.presto.ClientRequestFilterManager;
 import com.facebook.presto.ClientRequestFilterModule;
+import com.facebook.presto.builtin.tools.WorkerFunctionRegistryTool;
 import com.facebook.presto.dispatcher.QueryPrerequisitesManager;
 import com.facebook.presto.dispatcher.QueryPrerequisitesManagerModule;
 import com.facebook.presto.eventlistener.EventListenerManager;
@@ -43,6 +44,7 @@
 import com.facebook.presto.metadata.Catalog;
 import com.facebook.presto.metadata.CatalogManager;
 import com.facebook.presto.metadata.DiscoveryNodeManager;
+import com.facebook.presto.metadata.FunctionAndTypeManager;
 import com.facebook.presto.metadata.InternalNodeManager;
 import com.facebook.presto.metadata.SessionPropertyManager;
 import com.facebook.presto.metadata.StaticCatalogStore;
@@ -53,7 +55,10 @@
 import com.facebook.presto.security.AccessControlModule;
 import com.facebook.presto.server.security.PasswordAuthenticatorManager;
 import com.facebook.presto.server.security.PrestoAuthenticatorManager;
+import com.facebook.presto.server.security.SecurityConfig;
 import com.facebook.presto.server.security.ServerSecurityModule;
+import com.facebook.presto.server.security.oauth2.OAuth2Client;
+import com.facebook.presto.spi.function.SqlFunction;
 import com.facebook.presto.sql.analyzer.FeaturesConfig;
 import com.facebook.presto.sql.expressions.ExpressionOptimizerManager;
 import com.facebook.presto.sql.parser.SqlParserOptions;
@@ -85,6 +90,7 @@
 import static com.facebook.airlift.json.JsonBinder.jsonBinder;
 import static com.facebook.presto.server.PrestoSystemRequirements.verifyJvmRequirements;
 import static com.facebook.presto.server.PrestoSystemRequirements.verifySystemTimeIsReasonable;
+import static com.facebook.presto.server.security.SecurityConfig.AuthenticationType.OAUTH2;
 import static com.google.common.base.Strings.nullToEmpty;
 import static java.util.Objects.requireNonNull;
 
@@ -198,11 +204,24 @@ public void run()
             PluginNodeManager pluginNodeManager = new PluginNodeManager(nodeManager, nodeInfo.getEnvironment());
             planCheckerProviderManager.loadPlanCheckerProviders(pluginNodeManager);
 
+            if (injector.getInstance(FeaturesConfig.class).isBuiltInSidecarFunctionsEnabled()) {
+                List functions = injector.getInstance(WorkerFunctionRegistryTool.class).getWorkerFunctions();
+                injector.getInstance(FunctionAndTypeManager.class).registerWorkerFunctions(functions);
+            }
+
             injector.getInstance(ClientRequestFilterManager.class).loadClientRequestFilters();
             injector.getInstance(ExpressionOptimizerManager.class).loadExpressionOptimizerFactories();
 
+            injector.getInstance(FunctionAndTypeManager.class)
+                    .getBuiltInPluginFunctionNamespaceManager().triggerConflictCheckWithBuiltInFunctions();
+
             startAssociatedProcesses(injector);
 
+            SecurityConfig securityConfig = injector.getInstance(SecurityConfig.class);
+            if (securityConfig.getAuthenticationTypes().contains(OAUTH2)) {
+                injector.getInstance(OAuth2Client.class).load();
+            }
+
             injector.getInstance(Announcer.class).start();
 
             log.info("======== SERVER STARTED ========");
diff --git a/presto-main/src/main/java/com/facebook/presto/server/ServerInfoResource.java b/presto-main/src/main/java/com/facebook/presto/server/ServerInfoResource.java
index e24a08ab00559..8c686288caef8 100644
--- a/presto-main/src/main/java/com/facebook/presto/server/ServerInfoResource.java
+++ b/presto-main/src/main/java/com/facebook/presto/server/ServerInfoResource.java
@@ -19,6 +19,7 @@
 import com.facebook.presto.execution.resourceGroups.ResourceGroupManager;
 import com.facebook.presto.metadata.StaticCatalogStore;
 import com.facebook.presto.spi.NodeState;
+import com.facebook.presto.spi.NodeStats;
 import jakarta.annotation.security.RolesAllowed;
 import jakarta.inject.Inject;
 import jakarta.ws.rs.Consumes;
@@ -128,6 +129,16 @@ else if (!nodeResourceStatusProvider.hasResources() || !resourceGroupManager.isC
         }
     }
 
+    @GET
+    @Path("stats")
+    @Produces({APPLICATION_JSON, APPLICATION_THRIFT_BINARY, APPLICATION_THRIFT_COMPACT, APPLICATION_THRIFT_FB_COMPACT})
+    @RolesAllowed(ADMIN)
+    public NodeStats getServerStats()
+    {
+        NodeStats stats = new NodeStats(getServerState(), null);
+        return stats;
+    }
+
     @GET
     @Path("coordinator")
     @Produces(TEXT_PLAIN)
diff --git a/presto-main/src/main/java/com/facebook/presto/server/ServerMainModule.java b/presto-main/src/main/java/com/facebook/presto/server/ServerMainModule.java
index 634b8654634a6..f5d2543048a46 100644
--- a/presto-main/src/main/java/com/facebook/presto/server/ServerMainModule.java
+++ b/presto-main/src/main/java/com/facebook/presto/server/ServerMainModule.java
@@ -16,7 +16,10 @@
 import com.facebook.airlift.concurrent.BoundedExecutor;
 import com.facebook.airlift.configuration.AbstractConfigurationAwareModule;
 import com.facebook.airlift.discovery.client.ServiceAnnouncement;
+import com.facebook.airlift.http.client.HttpClient;
 import com.facebook.airlift.http.server.TheServlet;
+import com.facebook.airlift.json.JsonCodec;
+import com.facebook.airlift.json.JsonCodecFactory;
 import com.facebook.airlift.json.JsonObjectMapperProvider;
 import com.facebook.airlift.stats.GcMonitor;
 import com.facebook.airlift.stats.JmxGcMonitor;
@@ -32,6 +35,10 @@
 import com.facebook.presto.PagesIndexPageSorter;
 import com.facebook.presto.SystemSessionProperties;
 import com.facebook.presto.block.BlockJsonSerde;
+import com.facebook.presto.builtin.tools.ForNativeFunctionRegistryInfo;
+import com.facebook.presto.builtin.tools.NativeSidecarFunctionRegistryTool;
+import com.facebook.presto.builtin.tools.NativeSidecarRegistryToolConfig;
+import com.facebook.presto.builtin.tools.WorkerFunctionRegistryTool;
 import com.facebook.presto.catalogserver.CatalogServerClient;
 import com.facebook.presto.catalogserver.RandomCatalogServerAddressSelector;
 import com.facebook.presto.catalogserver.RemoteMetadataManager;
@@ -77,8 +84,11 @@
 import com.facebook.presto.execution.scheduler.NodeSchedulerConfig;
 import com.facebook.presto.execution.scheduler.NodeSchedulerExporter;
 import com.facebook.presto.execution.scheduler.TableWriteInfo;
+import com.facebook.presto.execution.scheduler.clusterOverload.ClusterOverloadPolicyModule;
+import com.facebook.presto.execution.scheduler.clusterOverload.ClusterResourceChecker;
 import com.facebook.presto.execution.scheduler.nodeSelection.NodeSelectionStats;
 import com.facebook.presto.execution.scheduler.nodeSelection.SimpleTtlNodeSelectorConfig;
+import com.facebook.presto.functionNamespace.JsonBasedUdfFunctionMetadata;
 import com.facebook.presto.index.IndexManager;
 import com.facebook.presto.memory.LocalMemoryManager;
 import com.facebook.presto.memory.LocalMemoryManagerExporter;
@@ -243,6 +253,7 @@
 import com.google.inject.Module;
 import com.google.inject.Provides;
 import com.google.inject.Scopes;
+import com.google.inject.TypeLiteral;
 import com.google.inject.multibindings.MapBinder;
 import io.airlift.slice.Slice;
 import jakarta.annotation.PreDestroy;
@@ -251,6 +262,7 @@
 import jakarta.servlet.Servlet;
 
 import java.util.List;
+import java.util.Map;
 import java.util.Optional;
 import java.util.concurrent.ExecutorService;
 import java.util.concurrent.LinkedBlockingQueue;
@@ -266,6 +278,7 @@
 import static com.facebook.airlift.http.client.HttpClientBinder.httpClientBinder;
 import static com.facebook.airlift.jaxrs.JaxrsBinder.jaxrsBinder;
 import static com.facebook.airlift.json.JsonBinder.jsonBinder;
+import static com.facebook.airlift.json.JsonCodec.listJsonCodec;
 import static com.facebook.airlift.json.JsonCodecBinder.jsonCodecBinder;
 import static com.facebook.airlift.json.smile.SmileCodecBinder.smileCodecBinder;
 import static com.facebook.airlift.units.DataSize.Unit.MEGABYTE;
@@ -394,6 +407,11 @@ else if (serverConfig.isCoordinator()) {
                 .withAddressSelector(((addressSelectorBinder, annotation, prefix) ->
                         addressSelectorBinder.bind(AddressSelector.class).annotatedWith(annotation).to(FixedAddressSelector.class)));
 
+        binder.bind(new TypeLiteral>>>() {})
+                .toInstance(new JsonCodecFactory().mapJsonCodec(String.class, listJsonCodec(JsonBasedUdfFunctionMetadata.class)));
+        httpClientBinder(binder).bindHttpClient("native-function-registry", ForNativeFunctionRegistryInfo.class);
+        configBinder(binder).bindConfig(NativeSidecarRegistryToolConfig.class);
+
         // node scheduler
         // TODO: remove from NodePartitioningManager and move to CoordinatorModule
         configBinder(binder).bindConfig(NodeSchedulerConfig.class);
@@ -693,6 +711,10 @@ public ListeningExecutorService createResourceManagerExecutor(ResourceManagerCon
         // system connector
         binder.install(new SystemConnectorModule());
 
+        // ClusterOverload policy module
+        binder.install(new ClusterOverloadPolicyModule());
+        newExporter(binder).export(ClusterResourceChecker.class).withGeneratedName();
+
         // splits
         jsonCodecBinder(binder).bindJsonCodec(TaskUpdateRequest.class);
         jsonCodecBinder(binder).bindJsonCodec(ConnectorSplit.class);
@@ -891,6 +913,22 @@ public static FragmentResultCacheManager createFragmentResultCacheManager(FileFr
         return new NoOpFragmentResultCacheManager();
     }
 
+    @Provides
+    @Singleton
+    public WorkerFunctionRegistryTool provideWorkerFunctionRegistryTool(
+            NativeSidecarRegistryToolConfig config,
+            @ForNativeFunctionRegistryInfo HttpClient httpClient,
+            JsonCodec>> nativeFunctionSignatureMapJsonCodec,
+            NodeManager nodeManager)
+    {
+        return new NativeSidecarFunctionRegistryTool(
+                httpClient,
+                nativeFunctionSignatureMapJsonCodec,
+                nodeManager,
+                config.getNativeSidecarRegistryToolNumRetries(),
+                config.getNativeSidecarRegistryToolRetryDelayMs());
+    }
+
     public static class ExecutorCleanup
     {
         private final List executors;
diff --git a/presto-main/src/main/java/com/facebook/presto/server/WebUiResource.java b/presto-main/src/main/java/com/facebook/presto/server/WebUiResource.java
index 2eeb67b5f5d0b..a2e908d47aea7 100644
--- a/presto-main/src/main/java/com/facebook/presto/server/WebUiResource.java
+++ b/presto-main/src/main/java/com/facebook/presto/server/WebUiResource.java
@@ -13,6 +13,7 @@
  */
 package com.facebook.presto.server;
 
+import com.facebook.presto.server.security.oauth2.OAuthWebUiCookie;
 import jakarta.annotation.security.RolesAllowed;
 import jakarta.ws.rs.GET;
 import jakarta.ws.rs.HeaderParam;
@@ -21,15 +22,19 @@
 import jakarta.ws.rs.core.Response;
 import jakarta.ws.rs.core.UriInfo;
 
+import java.util.Optional;
+
 import static com.facebook.presto.server.security.RoleType.ADMIN;
+import static com.facebook.presto.server.security.oauth2.OAuth2Utils.getLastURLParameter;
 import static com.google.common.base.Strings.isNullOrEmpty;
 import static com.google.common.net.HttpHeaders.X_FORWARDED_PROTO;
-import static jakarta.ws.rs.core.Response.Status.MOVED_PERMANENTLY;
 
 @Path("/")
 @RolesAllowed(ADMIN)
 public class WebUiResource
 {
+    public static final String UI_ENDPOINT = "/";
+
     @GET
     public Response redirectIndexHtml(
             @HeaderParam(X_FORWARDED_PROTO) String proto,
@@ -38,9 +43,30 @@ public Response redirectIndexHtml(
         if (isNullOrEmpty(proto)) {
             proto = uriInfo.getRequestUri().getScheme();
         }
+        Optional lastURL = getLastURLParameter(uriInfo.getQueryParameters());
+        if (lastURL.isPresent()) {
+            return Response
+                    .seeOther(uriInfo.getRequestUriBuilder().scheme(proto).uri(lastURL.get()).build())
+                    .build();
+        }
 
-        return Response.status(MOVED_PERMANENTLY)
-                .location(uriInfo.getRequestUriBuilder().scheme(proto).path("/ui/").build())
+        return Response
+                .temporaryRedirect(uriInfo.getRequestUriBuilder().scheme(proto).path("/ui/").replaceQuery("").build())
+                .build();
+    }
+
+    @GET
+    @Path("/logout")
+    public Response logout(
+            @HeaderParam(X_FORWARDED_PROTO) String proto,
+            @Context UriInfo uriInfo)
+    {
+        if (isNullOrEmpty(proto)) {
+            proto = uriInfo.getRequestUri().getScheme();
+        }
+        return Response
+                .temporaryRedirect(uriInfo.getBaseUriBuilder().scheme(proto).path("/ui/logout.html").build())
+                .cookie(OAuthWebUiCookie.delete())
                 .build();
     }
 }
diff --git a/presto-main/src/main/java/com/facebook/presto/server/security/AuthenticationFilter.java b/presto-main/src/main/java/com/facebook/presto/server/security/AuthenticationFilter.java
index 6e5907735c8a8..17d462cb1b822 100644
--- a/presto-main/src/main/java/com/facebook/presto/server/security/AuthenticationFilter.java
+++ b/presto-main/src/main/java/com/facebook/presto/server/security/AuthenticationFilter.java
@@ -16,6 +16,7 @@
 import com.facebook.airlift.http.server.AuthenticationException;
 import com.facebook.airlift.http.server.Authenticator;
 import com.facebook.presto.ClientRequestFilterManager;
+import com.facebook.presto.server.security.oauth2.OAuth2Authenticator;
 import com.facebook.presto.spi.ClientRequestFilter;
 import com.facebook.presto.spi.PrestoException;
 import com.google.common.base.Joiner;
@@ -45,7 +46,11 @@
 import java.util.List;
 import java.util.Map;
 import java.util.Set;
+import java.util.stream.Collectors;
 
+import static com.facebook.presto.server.WebUiResource.UI_ENDPOINT;
+import static com.facebook.presto.server.security.oauth2.OAuth2CallbackResource.CALLBACK_ENDPOINT;
+import static com.facebook.presto.server.security.oauth2.OAuth2TokenExchangeResource.TOKEN_ENDPOINT;
 import static com.facebook.presto.spi.StandardErrorCode.HEADER_MODIFICATION_ATTEMPT;
 import static com.google.common.io.ByteStreams.copy;
 import static com.google.common.io.ByteStreams.nullOutputStream;
@@ -64,15 +69,39 @@ public class AuthenticationFilter
     private final boolean allowForwardedHttps;
     private final ClientRequestFilterManager clientRequestFilterManager;
     private final List headersBlockList = ImmutableList.of("X-Presto-Transaction-Id", "X-Presto-Started-Transaction-Id", "X-Presto-Clear-Transaction-Id", "X-Presto-Trace-Token");
+    private final WebUiAuthenticationManager webUiAuthenticationManager;
+    private final boolean isOauth2Enabled;
 
     @Inject
-    public AuthenticationFilter(List authenticators, SecurityConfig securityConfig, ClientRequestFilterManager clientRequestFilterManager)
+    public AuthenticationFilter(List authenticators, SecurityConfig securityConfig, ClientRequestFilterManager clientRequestFilterManager, WebUiAuthenticationManager webUiAuthenticationManager)
     {
         this.authenticators = ImmutableList.copyOf(requireNonNull(authenticators, "authenticators is null"));
+        this.webUiAuthenticationManager = requireNonNull(webUiAuthenticationManager, "webUiAuthenticationManager is null");
+        this.isOauth2Enabled = this.authenticators.stream()
+                .anyMatch(a -> a.getClass().equals(OAuth2Authenticator.class));
         this.allowForwardedHttps = requireNonNull(securityConfig, "securityConfig is null").getAllowForwardedHttps();
         this.clientRequestFilterManager = requireNonNull(clientRequestFilterManager, "clientRequestFilterManager is null");
     }
 
+    public static ServletRequest withPrincipal(HttpServletRequest request, Principal principal)
+    {
+        requireNonNull(principal, "principal is null");
+        return new HttpServletRequestWrapper(request)
+        {
+            @Override
+            public Principal getUserPrincipal()
+            {
+                return principal;
+            }
+        };
+    }
+
+    private static boolean isRequestToOAuthEndpoint(HttpServletRequest request)
+    {
+        return request.getPathInfo().startsWith(TOKEN_ENDPOINT)
+                || request.getPathInfo().startsWith(CALLBACK_ENDPOINT);
+    }
+
     @Override
     public void init(FilterConfig filterConfig) {}
 
@@ -86,6 +115,13 @@ public void doFilter(ServletRequest servletRequest, ServletResponse servletRespo
         HttpServletRequest request = (HttpServletRequest) servletRequest;
         HttpServletResponse response = (HttpServletResponse) servletResponse;
 
+        // Check if it's a request going to the web UI side.
+        if (isWebUiRequest(request) && isOauth2Enabled) {
+            // call web authenticator
+            this.webUiAuthenticationManager.handleRequest(request, response, nextFilter);
+            return;
+        }
+
         // skip authentication if non-secure or not configured
         if (!doesRequestSupportAuthentication(request)) {
             nextFilter.doFilter(request, response);
@@ -108,6 +144,7 @@ public void doFilter(ServletRequest servletRequest, ServletResponse servletRespo
                 e.getAuthenticateHeader().ifPresent(authenticateHeaders::add);
                 continue;
             }
+
             // authentication succeeded
             HttpServletRequest wrappedRequest = mergeExtraHeaders(request, principal);
             nextFilter.doFilter(withPrincipal(wrappedRequest, principal), response);
@@ -117,6 +154,11 @@ public void doFilter(ServletRequest servletRequest, ServletResponse servletRespo
         // authentication failed
         skipRequestBody(request);
 
+        // Browsers have special handling for the BASIC challenge authenticate header so we need to filter them out if the WebUI Oauth Token is present.
+        if (isOauth2Enabled && OAuth2Authenticator.extractTokenFromCookie(request).isPresent()) {
+            authenticateHeaders = authenticateHeaders.stream().filter(value -> value.contains("x_token_server")).collect(Collectors.toSet());
+        }
+
         for (String value : authenticateHeaders) {
             response.addHeader(WWW_AUTHENTICATE, value);
         }
@@ -182,6 +224,9 @@ public HttpServletRequest mergeExtraHeaders(HttpServletRequest request, Principa
 
     private boolean doesRequestSupportAuthentication(HttpServletRequest request)
     {
+        if (isRequestToOAuthEndpoint(request)) {
+            return false;
+        }
         if (authenticators.isEmpty()) {
             return false;
         }
@@ -194,19 +239,6 @@ private boolean doesRequestSupportAuthentication(HttpServletRequest request)
         return false;
     }
 
-    private static ServletRequest withPrincipal(HttpServletRequest request, Principal principal)
-    {
-        requireNonNull(principal, "principal is null");
-        return new HttpServletRequestWrapper(request)
-        {
-            @Override
-            public Principal getUserPrincipal()
-            {
-                return principal;
-            }
-        };
-    }
-
     private static void skipRequestBody(HttpServletRequest request)
             throws IOException
     {
@@ -221,6 +253,12 @@ private static void skipRequestBody(HttpServletRequest request)
         }
     }
 
+    private boolean isWebUiRequest(HttpServletRequest request)
+    {
+        String pathInfo = request.getPathInfo();
+        return pathInfo == null || pathInfo.equals(UI_ENDPOINT) || pathInfo.startsWith("/ui");
+    }
+
     public static class ModifiedHttpServletRequest
             extends HttpServletRequestWrapper
     {
diff --git a/presto-main/src/main/java/com/facebook/presto/server/security/DefaultWebUiAuthenticationManager.java b/presto-main/src/main/java/com/facebook/presto/server/security/DefaultWebUiAuthenticationManager.java
new file mode 100644
index 0000000000000..75ff4d072d711
--- /dev/null
+++ b/presto-main/src/main/java/com/facebook/presto/server/security/DefaultWebUiAuthenticationManager.java
@@ -0,0 +1,32 @@
+/*
+ * Licensed 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 com.facebook.presto.server.security;
+
+import jakarta.servlet.FilterChain;
+import jakarta.servlet.ServletException;
+import jakarta.servlet.http.HttpServletRequest;
+import jakarta.servlet.http.HttpServletResponse;
+
+import java.io.IOException;
+
+public class DefaultWebUiAuthenticationManager
+        implements WebUiAuthenticationManager
+{
+    @Override
+    public void handleRequest(HttpServletRequest request, HttpServletResponse response, FilterChain nextFilter)
+            throws IOException, ServletException
+    {
+        nextFilter.doFilter(request, response);
+    }
+}
diff --git a/presto-main/src/main/java/com/facebook/presto/server/security/ServerSecurityModule.java b/presto-main/src/main/java/com/facebook/presto/server/security/ServerSecurityModule.java
index cd7943e40b313..6fe27fa7b7e96 100644
--- a/presto-main/src/main/java/com/facebook/presto/server/security/ServerSecurityModule.java
+++ b/presto-main/src/main/java/com/facebook/presto/server/security/ServerSecurityModule.java
@@ -15,13 +15,22 @@
 
 import com.facebook.airlift.configuration.AbstractConfigurationAwareModule;
 import com.facebook.airlift.http.server.Authenticator;
+import com.facebook.airlift.http.server.Authorizer;
 import com.facebook.airlift.http.server.CertificateAuthenticator;
+import com.facebook.airlift.http.server.ConfigurationBasedAuthorizer;
+import com.facebook.airlift.http.server.ConfigurationBasedAuthorizerConfig;
 import com.facebook.airlift.http.server.KerberosAuthenticator;
 import com.facebook.airlift.http.server.KerberosConfig;
+import com.facebook.airlift.http.server.TheServlet;
 import com.facebook.presto.server.security.SecurityConfig.AuthenticationType;
+import com.facebook.presto.server.security.oauth2.OAuth2AuthenticationSupportModule;
+import com.facebook.presto.server.security.oauth2.OAuth2Authenticator;
+import com.facebook.presto.server.security.oauth2.OAuth2Config;
+import com.facebook.presto.server.security.oauth2.OAuth2WebUiAuthenticationManager;
 import com.google.inject.Binder;
 import com.google.inject.Scopes;
 import com.google.inject.multibindings.Multibinder;
+import jakarta.servlet.Filter;
 
 import java.util.List;
 
@@ -30,8 +39,11 @@
 import static com.facebook.presto.server.security.SecurityConfig.AuthenticationType.CUSTOM;
 import static com.facebook.presto.server.security.SecurityConfig.AuthenticationType.JWT;
 import static com.facebook.presto.server.security.SecurityConfig.AuthenticationType.KERBEROS;
+import static com.facebook.presto.server.security.SecurityConfig.AuthenticationType.OAUTH2;
 import static com.facebook.presto.server.security.SecurityConfig.AuthenticationType.PASSWORD;
+import static com.facebook.presto.server.security.SecurityConfig.AuthenticationType.TEST_EXTERNAL;
 import static com.google.inject.multibindings.Multibinder.newSetBinder;
+import static com.google.inject.multibindings.OptionalBinder.newOptionalBinder;
 
 public class ServerSecurityModule
         extends AbstractConfigurationAwareModule
@@ -39,6 +51,10 @@ public class ServerSecurityModule
     @Override
     protected void setup(Binder binder)
     {
+        newOptionalBinder(binder, WebUiAuthenticationManager.class).setDefault().to(DefaultWebUiAuthenticationManager.class).in(Scopes.SINGLETON);
+        newSetBinder(binder, Filter.class, TheServlet.class).addBinding()
+                .to(AuthenticationFilter.class).in(Scopes.SINGLETON);
+
         binder.bind(PasswordAuthenticatorManager.class).in(Scopes.SINGLETON);
         binder.bind(PrestoAuthenticatorManager.class).in(Scopes.SINGLETON);
 
@@ -60,11 +76,25 @@ else if (authType == JWT) {
                 configBinder(binder).bindConfig(JsonWebTokenConfig.class);
                 authBinder.addBinding().to(JsonWebTokenAuthenticator.class).in(Scopes.SINGLETON);
             }
+            else if (authType == OAUTH2) {
+                newOptionalBinder(binder, WebUiAuthenticationManager.class).setBinding().to(OAuth2WebUiAuthenticationManager.class).in(Scopes.SINGLETON);
+                install(new OAuth2AuthenticationSupportModule());
+                binder.bind(OAuth2Authenticator.class).in(Scopes.SINGLETON);
+                configBinder(binder).bindConfig(OAuth2Config.class);
+                authBinder.addBinding().to(OAuth2Authenticator.class).in(Scopes.SINGLETON);
+
+                configBinder(binder).bindConfig(ConfigurationBasedAuthorizerConfig.class);
+                binder.bind(Authorizer.class).to(ConfigurationBasedAuthorizer.class).in(Scopes.SINGLETON);
+            }
             else if (authType == CUSTOM) {
                 authBinder.addBinding().to(CustomPrestoAuthenticator.class).in(Scopes.SINGLETON);
             }
             else {
-                throw new AssertionError("Unhandled auth type: " + authType);
+                // TEST_EXTERNAL is an authentication type used for testing the external auth flow for the JDBC driver.
+                // This is here as a guard since it's not a real authenticator but if I exclude it from the checks then teh error is thrown.
+                if (authType != TEST_EXTERNAL) {
+                    throw new AssertionError("Unhandled auth type: " + authType);
+                }
             }
         }
     }
diff --git a/presto-main/src/main/java/com/facebook/presto/server/security/WebUiAuthenticationManager.java b/presto-main/src/main/java/com/facebook/presto/server/security/WebUiAuthenticationManager.java
new file mode 100644
index 0000000000000..322d416168942
--- /dev/null
+++ b/presto-main/src/main/java/com/facebook/presto/server/security/WebUiAuthenticationManager.java
@@ -0,0 +1,27 @@
+/*
+ * Licensed 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 com.facebook.presto.server.security;
+
+import jakarta.servlet.FilterChain;
+import jakarta.servlet.ServletException;
+import jakarta.servlet.http.HttpServletRequest;
+import jakarta.servlet.http.HttpServletResponse;
+
+import java.io.IOException;
+
+public interface WebUiAuthenticationManager
+{
+    void handleRequest(HttpServletRequest request, HttpServletResponse response, FilterChain nextFilter)
+            throws IOException, ServletException;
+}
diff --git a/presto-main/src/main/java/com/facebook/presto/server/security/oauth2/ChallengeFailedException.java b/presto-main/src/main/java/com/facebook/presto/server/security/oauth2/ChallengeFailedException.java
new file mode 100644
index 0000000000000..f72a7c7c537ea
--- /dev/null
+++ b/presto-main/src/main/java/com/facebook/presto/server/security/oauth2/ChallengeFailedException.java
@@ -0,0 +1,28 @@
+/*
+ * Licensed 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 com.facebook.presto.server.security.oauth2;
+
+public class ChallengeFailedException
+        extends Exception
+{
+    public ChallengeFailedException(String message)
+    {
+        super(message);
+    }
+
+    public ChallengeFailedException(String message, Throwable cause)
+    {
+        super(message, cause);
+    }
+}
diff --git a/presto-main/src/main/java/com/facebook/presto/server/security/oauth2/ForOAuth2.java b/presto-main/src/main/java/com/facebook/presto/server/security/oauth2/ForOAuth2.java
new file mode 100644
index 0000000000000..ee66543301512
--- /dev/null
+++ b/presto-main/src/main/java/com/facebook/presto/server/security/oauth2/ForOAuth2.java
@@ -0,0 +1,31 @@
+/*
+ * Licensed 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 com.facebook.presto.server.security.oauth2;
+
+import com.google.inject.BindingAnnotation;
+
+import java.lang.annotation.Retention;
+import java.lang.annotation.Target;
+
+import static java.lang.annotation.ElementType.FIELD;
+import static java.lang.annotation.ElementType.METHOD;
+import static java.lang.annotation.ElementType.PARAMETER;
+import static java.lang.annotation.RetentionPolicy.RUNTIME;
+
+@Retention(RUNTIME)
+@Target({FIELD, PARAMETER, METHOD})
+@BindingAnnotation
+public @interface ForOAuth2
+{
+}
diff --git a/presto-main/src/main/java/com/facebook/presto/server/security/oauth2/ForRefreshTokens.java b/presto-main/src/main/java/com/facebook/presto/server/security/oauth2/ForRefreshTokens.java
new file mode 100644
index 0000000000000..d4d3939c61320
--- /dev/null
+++ b/presto-main/src/main/java/com/facebook/presto/server/security/oauth2/ForRefreshTokens.java
@@ -0,0 +1,31 @@
+/*
+ * Licensed 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 com.facebook.presto.server.security.oauth2;
+
+import com.google.inject.BindingAnnotation;
+
+import java.lang.annotation.Retention;
+import java.lang.annotation.Target;
+
+import static java.lang.annotation.ElementType.FIELD;
+import static java.lang.annotation.ElementType.METHOD;
+import static java.lang.annotation.ElementType.PARAMETER;
+import static java.lang.annotation.RetentionPolicy.RUNTIME;
+
+@Retention(RUNTIME)
+@Target({FIELD, PARAMETER, METHOD})
+@BindingAnnotation
+public @interface ForRefreshTokens
+{
+}
diff --git a/presto-main/src/main/java/com/facebook/presto/server/security/oauth2/JweTokenSerializer.java b/presto-main/src/main/java/com/facebook/presto/server/security/oauth2/JweTokenSerializer.java
new file mode 100644
index 0000000000000..b10ea12a8e8a2
--- /dev/null
+++ b/presto-main/src/main/java/com/facebook/presto/server/security/oauth2/JweTokenSerializer.java
@@ -0,0 +1,179 @@
+/*
+ * Licensed 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 com.facebook.presto.server.security.oauth2;
+
+import com.facebook.airlift.units.Duration;
+import com.nimbusds.jose.EncryptionMethod;
+import com.nimbusds.jose.JOSEException;
+import com.nimbusds.jose.JWEAlgorithm;
+import com.nimbusds.jose.JWEHeader;
+import com.nimbusds.jose.JWEObject;
+import com.nimbusds.jose.KeyLengthException;
+import com.nimbusds.jose.Payload;
+import com.nimbusds.jose.crypto.AESDecrypter;
+import com.nimbusds.jose.crypto.AESEncrypter;
+import io.jsonwebtoken.Claims;
+import io.jsonwebtoken.CompressionCodec;
+import io.jsonwebtoken.CompressionException;
+import io.jsonwebtoken.Header;
+import io.jsonwebtoken.JwtBuilder;
+import io.jsonwebtoken.JwtParser;
+
+import javax.crypto.KeyGenerator;
+import javax.crypto.SecretKey;
+
+import java.security.NoSuchAlgorithmException;
+import java.text.ParseException;
+import java.time.Clock;
+import java.util.Date;
+import java.util.Map;
+import java.util.Optional;
+
+import static com.facebook.presto.server.security.oauth2.JwtUtil.newJwtBuilder;
+import static com.facebook.presto.server.security.oauth2.JwtUtil.newJwtParserBuilder;
+import static com.google.common.base.Preconditions.checkState;
+import static java.lang.String.format;
+import static java.util.Objects.requireNonNull;
+
+public class JweTokenSerializer
+        implements TokenPairSerializer
+{
+    private static final JWEAlgorithm ALGORITHM = JWEAlgorithm.A256KW;
+    private static final EncryptionMethod ENCRYPTION_METHOD = EncryptionMethod.A256CBC_HS512;
+    private static final CompressionCodec COMPRESSION_CODEC = new ZstdCodec();
+    private static final String ACCESS_TOKEN_KEY = "access_token";
+    private static final String EXPIRATION_TIME_KEY = "expiration_time";
+    private static final String REFRESH_TOKEN_KEY = "refresh_token";
+    private final OAuth2Client client;
+    private final Clock clock;
+    private final String issuer;
+    private final String audience;
+    private final Duration tokenExpiration;
+    private final JwtParser parser;
+    private final AESEncrypter jweEncrypter;
+    private final AESDecrypter jweDecrypter;
+    private final String principalField;
+
+    public JweTokenSerializer(
+            RefreshTokensConfig config,
+            OAuth2Client client,
+            String issuer,
+            String audience,
+            String principalField,
+            Clock clock,
+            Duration tokenExpiration)
+            throws KeyLengthException, NoSuchAlgorithmException
+    {
+        SecretKey secretKey = createKey(requireNonNull(config, "config is null"));
+        this.jweEncrypter = new AESEncrypter(secretKey);
+        this.jweDecrypter = new AESDecrypter(secretKey);
+        this.client = requireNonNull(client, "client is null");
+        this.issuer = requireNonNull(issuer, "issuer is null");
+        this.principalField = requireNonNull(principalField, "principalField is null");
+        this.audience = requireNonNull(audience, "issuer is null");
+        this.clock = requireNonNull(clock, "clock is null");
+        this.tokenExpiration = requireNonNull(tokenExpiration, "tokenExpiration is null");
+
+        this.parser = newJwtParserBuilder()
+                .setClock(() -> Date.from(clock.instant()))
+                .requireIssuer(this.issuer)
+                .requireAudience(this.audience)
+                .setCompressionCodecResolver(JweTokenSerializer::resolveCompressionCodec)
+                .build();
+    }
+
+    @Override
+    public TokenPair deserialize(String token)
+    {
+        requireNonNull(token, "token is null");
+
+        try {
+            JWEObject jwe = JWEObject.parse(token);
+            jwe.decrypt(jweDecrypter);
+            Claims claims = parser.parseClaimsJwt(jwe.getPayload().toString()).getBody();
+            return TokenPair.accessAndRefreshTokens(
+                    claims.get(ACCESS_TOKEN_KEY, String.class),
+                    claims.get(EXPIRATION_TIME_KEY, Date.class),
+                    claims.get(REFRESH_TOKEN_KEY, String.class));
+        }
+        catch (ParseException ex) {
+            throw new IllegalArgumentException("Malformed jwt token", ex);
+        }
+        catch (JOSEException ex) {
+            throw new IllegalArgumentException("Decryption failed", ex);
+        }
+    }
+
+    @Override
+    public String serialize(TokenPair tokenPair)
+    {
+        requireNonNull(tokenPair, "tokenPair is null");
+
+        Optional> accessTokenClaims = client.getClaims(tokenPair.getAccessToken());
+        if (!accessTokenClaims.isPresent()) {
+            throw new IllegalArgumentException("Claims are missing");
+        }
+        Map claims = accessTokenClaims.get();
+        if (!claims.containsKey(principalField)) {
+            throw new IllegalArgumentException(format("%s field is missing", principalField));
+        }
+        JwtBuilder jwt = newJwtBuilder()
+                .setExpiration(Date.from(clock.instant().plusMillis(tokenExpiration.toMillis())))
+                .claim(principalField, claims.get(principalField).toString())
+                .setAudience(audience)
+                .setIssuer(issuer)
+                .claim(ACCESS_TOKEN_KEY, tokenPair.getAccessToken())
+                .claim(EXPIRATION_TIME_KEY, tokenPair.getExpiration())
+                .claim(REFRESH_TOKEN_KEY, tokenPair.getRefreshToken().orElseThrow(JweTokenSerializer::throwExceptionForNonExistingRefreshToken))
+                .compressWith(COMPRESSION_CODEC);
+
+        try {
+            JWEObject jwe = new JWEObject(
+                    new JWEHeader(ALGORITHM, ENCRYPTION_METHOD),
+                    new Payload(jwt.compact()));
+            jwe.encrypt(jweEncrypter);
+            return jwe.serialize();
+        }
+        catch (JOSEException ex) {
+            throw new IllegalStateException("Encryption failed", ex);
+        }
+    }
+
+    private static SecretKey createKey(RefreshTokensConfig config)
+            throws NoSuchAlgorithmException
+    {
+        SecretKey signingKey = config.getSecretKey();
+        if (signingKey == null) {
+            KeyGenerator generator = KeyGenerator.getInstance("AES");
+            generator.init(256);
+            return generator.generateKey();
+        }
+        return signingKey;
+    }
+
+    private static RuntimeException throwExceptionForNonExistingRefreshToken()
+    {
+        throw new IllegalStateException("Expected refresh token to be present. Please check your identity provider setup, or disable refresh tokens");
+    }
+
+    private static CompressionCodec resolveCompressionCodec(Header header)
+            throws CompressionException
+    {
+        if (header.getCompressionAlgorithm() != null) {
+            checkState(header.getCompressionAlgorithm().equals(ZstdCodec.CODEC_NAME), "Unknown codec '%s' used for token compression", header.getCompressionAlgorithm());
+            return COMPRESSION_CODEC;
+        }
+        return null;
+    }
+}
diff --git a/presto-main/src/main/java/com/facebook/presto/server/security/oauth2/JweTokenSerializerModule.java b/presto-main/src/main/java/com/facebook/presto/server/security/oauth2/JweTokenSerializerModule.java
new file mode 100644
index 0000000000000..1819f79a33137
--- /dev/null
+++ b/presto-main/src/main/java/com/facebook/presto/server/security/oauth2/JweTokenSerializerModule.java
@@ -0,0 +1,62 @@
+/*
+ * Licensed 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 com.facebook.presto.server.security.oauth2;
+
+import com.facebook.airlift.configuration.AbstractConfigurationAwareModule;
+import com.facebook.presto.client.NodeVersion;
+import com.google.inject.Binder;
+import com.google.inject.Inject;
+import com.google.inject.Key;
+import com.google.inject.Provides;
+import com.google.inject.Singleton;
+import com.nimbusds.jose.KeyLengthException;
+
+import java.security.NoSuchAlgorithmException;
+import java.time.Clock;
+import java.time.Duration;
+
+import static com.facebook.airlift.configuration.ConfigBinder.configBinder;
+import static com.google.inject.multibindings.OptionalBinder.newOptionalBinder;
+
+public class JweTokenSerializerModule
+        extends AbstractConfigurationAwareModule
+{
+    @Override
+    protected void setup(Binder binder)
+    {
+        configBinder(binder).bindConfig(RefreshTokensConfig.class);
+        RefreshTokensConfig config = buildConfigObject(RefreshTokensConfig.class);
+        newOptionalBinder(binder, Key.get(Duration.class, ForRefreshTokens.class)).setBinding().toInstance(Duration.ofMillis(config.getTokenExpiration().toMillis()));
+    }
+
+    @Provides
+    @Singleton
+    @Inject
+    public TokenPairSerializer getTokenPairSerializer(
+            OAuth2Client client,
+            NodeVersion nodeVersion,
+            RefreshTokensConfig config,
+            OAuth2Config oAuth2Config)
+            throws KeyLengthException, NoSuchAlgorithmException
+    {
+        return new JweTokenSerializer(
+                config,
+                client,
+                config.getIssuer() + "_" + nodeVersion.getVersion(),
+                config.getAudience(),
+                oAuth2Config.getPrincipalField(),
+                Clock.systemUTC(),
+                config.getTokenExpiration());
+    }
+}
diff --git a/presto-main/src/main/java/com/facebook/presto/server/security/oauth2/JwtUtil.java b/presto-main/src/main/java/com/facebook/presto/server/security/oauth2/JwtUtil.java
new file mode 100644
index 0000000000000..c2ea7114fd322
--- /dev/null
+++ b/presto-main/src/main/java/com/facebook/presto/server/security/oauth2/JwtUtil.java
@@ -0,0 +1,45 @@
+/*
+ * Licensed 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 com.facebook.presto.server.security.oauth2;
+
+import io.jsonwebtoken.JwtBuilder;
+import io.jsonwebtoken.JwtParserBuilder;
+import io.jsonwebtoken.Jwts;
+import io.jsonwebtoken.io.Deserializer;
+import io.jsonwebtoken.io.Serializer;
+import io.jsonwebtoken.jackson.io.JacksonDeserializer;
+import io.jsonwebtoken.jackson.io.JacksonSerializer;
+
+import java.util.Map;
+
+// avoid reflection and services lookup
+public final class JwtUtil
+{
+    private static final Serializer> JWT_SERIALIZER = new JacksonSerializer<>();
+    private static final Deserializer> JWT_DESERIALIZER = new JacksonDeserializer<>();
+
+    private JwtUtil() {}
+
+    public static JwtBuilder newJwtBuilder()
+    {
+        return Jwts.builder()
+                .serializeToJsonWith(JWT_SERIALIZER);
+    }
+
+    public static JwtParserBuilder newJwtParserBuilder()
+    {
+        return Jwts.parserBuilder()
+                .deserializeJsonWith(JWT_DESERIALIZER);
+    }
+}
diff --git a/presto-main/src/main/java/com/facebook/presto/server/security/oauth2/NimbusAirliftHttpClient.java b/presto-main/src/main/java/com/facebook/presto/server/security/oauth2/NimbusAirliftHttpClient.java
new file mode 100644
index 0000000000000..3aeeb4046f216
--- /dev/null
+++ b/presto-main/src/main/java/com/facebook/presto/server/security/oauth2/NimbusAirliftHttpClient.java
@@ -0,0 +1,137 @@
+/*
+ * Licensed 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 com.facebook.presto.server.security.oauth2;
+
+import com.facebook.airlift.http.client.HttpClient;
+import com.facebook.airlift.http.client.Request;
+import com.facebook.airlift.http.client.Response;
+import com.facebook.airlift.http.client.ResponseHandler;
+import com.facebook.airlift.http.client.ResponseHandlerUtils;
+import com.facebook.airlift.http.client.StringResponseHandler;
+import com.google.common.collect.ImmutableMultimap;
+import com.nimbusds.jose.util.Resource;
+import com.nimbusds.oauth2.sdk.ParseException;
+import com.nimbusds.oauth2.sdk.http.HTTPRequest;
+import com.nimbusds.oauth2.sdk.http.HTTPResponse;
+import jakarta.ws.rs.core.UriBuilder;
+
+import javax.inject.Inject;
+
+import java.io.IOException;
+import java.net.URISyntaxException;
+import java.net.URL;
+
+import static com.facebook.airlift.http.client.Request.Builder.prepareGet;
+import static com.facebook.airlift.http.client.StaticBodyGenerator.createStaticBodyGenerator;
+import static com.facebook.airlift.http.client.StringResponseHandler.createStringResponseHandler;
+import static com.google.common.net.HttpHeaders.CONTENT_TYPE;
+import static com.nimbusds.oauth2.sdk.http.HTTPRequest.Method.DELETE;
+import static com.nimbusds.oauth2.sdk.http.HTTPRequest.Method.GET;
+import static com.nimbusds.oauth2.sdk.http.HTTPRequest.Method.POST;
+import static com.nimbusds.oauth2.sdk.http.HTTPRequest.Method.PUT;
+import static java.lang.String.format;
+import static java.nio.charset.StandardCharsets.UTF_8;
+import static java.util.Objects.requireNonNull;
+
+public class NimbusAirliftHttpClient
+        implements NimbusHttpClient
+{
+    private final HttpClient httpClient;
+
+    @Inject
+    public NimbusAirliftHttpClient(@ForOAuth2 HttpClient httpClient)
+    {
+        this.httpClient = requireNonNull(httpClient, "httpClient is null");
+    }
+
+    @Override
+    public Resource retrieveResource(URL url)
+            throws IOException
+    {
+        try {
+            StringResponseHandler.StringResponse response = httpClient.execute(
+                    prepareGet().setUri(url.toURI()).build(),
+                    createStringResponseHandler());
+            return new Resource(response.getBody(), response.getHeader(CONTENT_TYPE));
+        }
+        catch (URISyntaxException e) {
+            throw new RuntimeException(e);
+        }
+    }
+
+    @Override
+    public  T execute(com.nimbusds.oauth2.sdk.Request nimbusRequest, Parser parser)
+    {
+        HTTPRequest httpRequest = nimbusRequest.toHTTPRequest();
+        HTTPRequest.Method method = httpRequest.getMethod();
+
+        Request.Builder request = new Request.Builder()
+                .setMethod(method.name())
+                .setFollowRedirects(httpRequest.getFollowRedirects());
+
+        UriBuilder url = UriBuilder.fromUri(httpRequest.getURI());
+        if (method.equals(GET) || method.equals(DELETE)) {
+            httpRequest.getQueryParameters().forEach((key, value) -> url.queryParam(key, value.toArray()));
+        }
+
+        url.fragment(httpRequest.getFragment());
+
+        request.setUri(url.build());
+
+        ImmutableMultimap.Builder headers = ImmutableMultimap.builder();
+        httpRequest.getHeaderMap().forEach(headers::putAll);
+        request.addHeaders(headers.build());
+
+        if (method.equals(POST) || method.equals(PUT)) {
+            String query = httpRequest.getQuery();
+            if (query != null) {
+                request.setBodyGenerator(createStaticBodyGenerator(httpRequest.getQuery(), UTF_8));
+            }
+        }
+        return httpClient.execute(request.build(), new NimbusResponseHandler<>(parser));
+    }
+
+    public static class NimbusResponseHandler
+            implements ResponseHandler
+    {
+        private final StringResponseHandler handler = createStringResponseHandler();
+        private final Parser parser;
+
+        public NimbusResponseHandler(Parser parser)
+        {
+            this.parser = requireNonNull(parser, "parser is null");
+        }
+
+        @Override
+        public T handleException(Request request, Exception exception)
+        {
+            throw ResponseHandlerUtils.propagate(request, exception);
+        }
+
+        @Override
+        public T handle(Request request, Response response)
+        {
+            StringResponseHandler.StringResponse stringResponse = handler.handle(request, response);
+            HTTPResponse nimbusResponse = new HTTPResponse(response.getStatusCode());
+            response.getHeaders().asMap().forEach((name, values) -> nimbusResponse.setHeader(name.toString(), values.toArray(new String[0])));
+            nimbusResponse.setContent(stringResponse.getBody());
+            try {
+                return parser.parse(nimbusResponse);
+            }
+            catch (ParseException e) {
+                throw new RuntimeException(format("Unable to parse response status=[%d], body=[%s]", stringResponse.getStatusCode(), stringResponse.getBody()), e);
+            }
+        }
+    }
+}
diff --git a/presto-main/src/main/java/com/facebook/presto/server/security/oauth2/NimbusHttpClient.java b/presto-main/src/main/java/com/facebook/presto/server/security/oauth2/NimbusHttpClient.java
new file mode 100644
index 0000000000000..b8c0e491dec4b
--- /dev/null
+++ b/presto-main/src/main/java/com/facebook/presto/server/security/oauth2/NimbusHttpClient.java
@@ -0,0 +1,31 @@
+/*
+ * Licensed 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 com.facebook.presto.server.security.oauth2;
+
+import com.nimbusds.jose.util.ResourceRetriever;
+import com.nimbusds.oauth2.sdk.ParseException;
+import com.nimbusds.oauth2.sdk.Request;
+import com.nimbusds.oauth2.sdk.http.HTTPResponse;
+
+public interface NimbusHttpClient
+        extends ResourceRetriever
+{
+     T execute(Request request, Parser parser);
+
+    interface Parser
+    {
+        T parse(HTTPResponse response)
+                throws ParseException;
+    }
+}
diff --git a/presto-main/src/main/java/com/facebook/presto/server/security/oauth2/NimbusOAuth2Client.java b/presto-main/src/main/java/com/facebook/presto/server/security/oauth2/NimbusOAuth2Client.java
new file mode 100644
index 0000000000000..86b2fb96b7d8f
--- /dev/null
+++ b/presto-main/src/main/java/com/facebook/presto/server/security/oauth2/NimbusOAuth2Client.java
@@ -0,0 +1,489 @@
+/*
+ * Licensed 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 com.facebook.presto.server.security.oauth2;
+
+import com.facebook.airlift.log.Logger;
+import com.facebook.airlift.units.Duration;
+import com.facebook.presto.server.security.oauth2.OAuth2ServerConfigProvider.OAuth2ServerConfig;
+import com.google.common.collect.ImmutableSet;
+import com.google.common.collect.Ordering;
+import com.nimbusds.jose.JOSEException;
+import com.nimbusds.jose.JWSAlgorithm;
+import com.nimbusds.jose.JWSHeader;
+import com.nimbusds.jose.jwk.source.RemoteJWKSet;
+import com.nimbusds.jose.proc.BadJOSEException;
+import com.nimbusds.jose.proc.JWSKeySelector;
+import com.nimbusds.jose.proc.JWSVerificationKeySelector;
+import com.nimbusds.jose.proc.SecurityContext;
+import com.nimbusds.jwt.JWTClaimsSet;
+import com.nimbusds.jwt.proc.DefaultJWTClaimsVerifier;
+import com.nimbusds.jwt.proc.DefaultJWTProcessor;
+import com.nimbusds.jwt.proc.JWTProcessor;
+import com.nimbusds.oauth2.sdk.AccessTokenResponse;
+import com.nimbusds.oauth2.sdk.AuthorizationCode;
+import com.nimbusds.oauth2.sdk.AuthorizationCodeGrant;
+import com.nimbusds.oauth2.sdk.AuthorizationGrant;
+import com.nimbusds.oauth2.sdk.AuthorizationRequest;
+import com.nimbusds.oauth2.sdk.ParseException;
+import com.nimbusds.oauth2.sdk.RefreshTokenGrant;
+import com.nimbusds.oauth2.sdk.Scope;
+import com.nimbusds.oauth2.sdk.TokenRequest;
+import com.nimbusds.oauth2.sdk.auth.ClientSecretBasic;
+import com.nimbusds.oauth2.sdk.auth.Secret;
+import com.nimbusds.oauth2.sdk.http.HTTPResponse;
+import com.nimbusds.oauth2.sdk.id.ClientID;
+import com.nimbusds.oauth2.sdk.id.Issuer;
+import com.nimbusds.oauth2.sdk.id.State;
+import com.nimbusds.oauth2.sdk.token.AccessToken;
+import com.nimbusds.oauth2.sdk.token.BearerAccessToken;
+import com.nimbusds.oauth2.sdk.token.RefreshToken;
+import com.nimbusds.oauth2.sdk.token.Tokens;
+import com.nimbusds.openid.connect.sdk.AuthenticationRequest;
+import com.nimbusds.openid.connect.sdk.Nonce;
+import com.nimbusds.openid.connect.sdk.OIDCTokenResponse;
+import com.nimbusds.openid.connect.sdk.UserInfoErrorResponse;
+import com.nimbusds.openid.connect.sdk.UserInfoRequest;
+import com.nimbusds.openid.connect.sdk.UserInfoResponse;
+import com.nimbusds.openid.connect.sdk.UserInfoSuccessResponse;
+import com.nimbusds.openid.connect.sdk.claims.AccessTokenHash;
+import com.nimbusds.openid.connect.sdk.claims.IDTokenClaimsSet;
+import com.nimbusds.openid.connect.sdk.token.OIDCTokens;
+import com.nimbusds.openid.connect.sdk.validators.AccessTokenValidator;
+import com.nimbusds.openid.connect.sdk.validators.IDTokenValidator;
+import com.nimbusds.openid.connect.sdk.validators.InvalidHashException;
+import net.minidev.json.JSONObject;
+
+import javax.inject.Inject;
+
+import java.net.MalformedURLException;
+import java.net.URI;
+import java.time.Instant;
+import java.util.Collections;
+import java.util.Date;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.Optional;
+import java.util.Set;
+import java.util.stream.Stream;
+
+import static com.google.common.base.Preconditions.checkArgument;
+import static com.google.common.base.Preconditions.checkState;
+import static com.google.common.collect.ImmutableList.toImmutableList;
+import static com.google.common.collect.ImmutableSet.toImmutableSet;
+import static com.google.common.hash.Hashing.sha256;
+import static com.nimbusds.oauth2.sdk.ResponseType.CODE;
+import static com.nimbusds.openid.connect.sdk.OIDCScopeValue.OPENID;
+import static java.nio.charset.StandardCharsets.UTF_8;
+import static java.util.Objects.requireNonNull;
+import static java.util.concurrent.TimeUnit.SECONDS;
+
+public class NimbusOAuth2Client
+        implements OAuth2Client
+{
+    private static final Logger LOG = Logger.get(NimbusAirliftHttpClient.class);
+
+    private final Issuer issuer;
+    private final ClientID clientId;
+    private final ClientSecretBasic clientAuth;
+    private final Scope scope;
+    private final String principalField;
+    private final Set accessTokenAudiences;
+    private final Duration maxClockSkew;
+    private final NimbusHttpClient httpClient;
+    private final OAuth2ServerConfigProvider serverConfigurationProvider;
+    private volatile boolean loaded;
+    private URI authUrl;
+    private URI tokenUrl;
+    private Optional userinfoUrl;
+    private JWSKeySelector jwsKeySelector;
+    private JWTProcessor accessTokenProcessor;
+    private AuthorizationCodeFlow flow;
+
+    @Inject
+    public NimbusOAuth2Client(OAuth2Config oauthConfig, OAuth2ServerConfigProvider serverConfigurationProvider, NimbusHttpClient httpClient)
+    {
+        requireNonNull(oauthConfig, "oauthConfig is null");
+        issuer = new Issuer(oauthConfig.getIssuer());
+        clientId = new ClientID(oauthConfig.getClientId());
+        clientAuth = new ClientSecretBasic(clientId, new Secret(oauthConfig.getClientSecret()));
+        scope = Scope.parse(oauthConfig.getScopes());
+        principalField = oauthConfig.getPrincipalField();
+        maxClockSkew = oauthConfig.getMaxClockSkew();
+
+        accessTokenAudiences = new HashSet<>(oauthConfig.getAdditionalAudiences());
+        accessTokenAudiences.add(clientId.getValue());
+        accessTokenAudiences.add(null); // A null value in the set allows JWTs with no audience
+
+        this.serverConfigurationProvider = requireNonNull(serverConfigurationProvider, "serverConfigurationProvider is null");
+        this.httpClient = requireNonNull(httpClient, "httpClient is null");
+    }
+
+    @Override
+    public void load()
+    {
+        OAuth2ServerConfig config = serverConfigurationProvider.get();
+        this.authUrl = config.getAuthUrl();
+        this.tokenUrl = config.getTokenUrl();
+        this.userinfoUrl = config.getUserinfoUrl();
+        try {
+            jwsKeySelector = new JWSVerificationKeySelector<>(
+                    Stream.concat(JWSAlgorithm.Family.RSA.stream(), JWSAlgorithm.Family.EC.stream()).collect(toImmutableSet()),
+                    new RemoteJWKSet<>(config.getJwksUrl().toURL(), httpClient));
+        }
+        catch (MalformedURLException e) {
+            throw new RuntimeException(e);
+        }
+
+        DefaultJWTProcessor processor = new DefaultJWTProcessor<>();
+        processor.setJWSKeySelector(jwsKeySelector);
+        DefaultJWTClaimsVerifier accessTokenVerifier = new DefaultJWTClaimsVerifier<>(
+                accessTokenAudiences,
+                new JWTClaimsSet.Builder()
+                        .issuer(config.getAccessTokenIssuer().orElse(issuer.getValue()))
+                        .build(),
+                ImmutableSet.of(principalField),
+                ImmutableSet.of());
+        accessTokenVerifier.setMaxClockSkew((int) maxClockSkew.roundTo(SECONDS));
+        processor.setJWTClaimsSetVerifier(accessTokenVerifier);
+        accessTokenProcessor = processor;
+        flow = scope.contains(OPENID) ? new OAuth2WithOidcExtensionsCodeFlow() : new OAuth2AuthorizationCodeFlow();
+        loaded = true;
+    }
+
+    @Override
+    public Request createAuthorizationRequest(String state, URI callbackUri)
+    {
+        checkState(loaded, "OAuth2 client not initialized");
+        return flow.createAuthorizationRequest(state, callbackUri);
+    }
+
+    @Override
+    public Response getOAuth2Response(String code, URI callbackUri, Optional nonce)
+            throws ChallengeFailedException
+    {
+        checkState(loaded, "OAuth2 client not initialized");
+        return flow.getOAuth2Response(code, callbackUri, nonce);
+    }
+
+    @Override
+    public Optional> getClaims(String accessToken)
+    {
+        checkState(loaded, "OAuth2 client not initialized");
+        return getJWTClaimsSet(accessToken).map(JWTClaimsSet::getClaims);
+    }
+
+    @Override
+    public Response refreshTokens(String refreshToken)
+            throws ChallengeFailedException
+    {
+        checkState(loaded, "OAuth2 client not initialized");
+        return flow.refreshTokens(refreshToken);
+    }
+
+    private interface AuthorizationCodeFlow
+    {
+        Request createAuthorizationRequest(String state, URI callbackUri);
+
+        Response getOAuth2Response(String code, URI callbackUri, Optional nonce)
+                throws ChallengeFailedException;
+
+        Response refreshTokens(String refreshToken)
+                throws ChallengeFailedException;
+    }
+
+    private class OAuth2AuthorizationCodeFlow
+            implements AuthorizationCodeFlow
+    {
+        @Override
+        public Request createAuthorizationRequest(String state, URI callbackUri)
+        {
+            return new Request(
+                    new AuthorizationRequest.Builder(CODE, clientId)
+                            .redirectionURI(callbackUri)
+                            .scope(scope)
+                            .endpointURI(authUrl)
+                            .state(new State(state))
+                            .build()
+                            .toURI(),
+                    Optional.empty());
+        }
+
+        @Override
+        public Response getOAuth2Response(String code, URI callbackUri, Optional nonce)
+                throws ChallengeFailedException
+        {
+            checkArgument(!nonce.isPresent(), "Unexpected nonce provided");
+            AccessTokenResponse tokenResponse = getTokenResponse(code, callbackUri, AccessTokenResponse::parse);
+            Tokens tokens = tokenResponse.toSuccessResponse().getTokens();
+            return toResponse(tokens, Optional.empty());
+        }
+
+        @Override
+        public Response refreshTokens(String refreshToken)
+                throws ChallengeFailedException
+        {
+            requireNonNull(refreshToken, "refreshToken is null");
+            AccessTokenResponse tokenResponse = getTokenResponse(refreshToken, AccessTokenResponse::parse);
+            return toResponse(tokenResponse.toSuccessResponse().getTokens(), Optional.of(refreshToken));
+        }
+
+        private Response toResponse(Tokens tokens, Optional existingRefreshToken)
+                throws ChallengeFailedException
+        {
+            AccessToken accessToken = tokens.getAccessToken();
+            RefreshToken refreshToken = tokens.getRefreshToken();
+            JWTClaimsSet claims = getJWTClaimsSet(accessToken.getValue()).orElseThrow(() -> new ChallengeFailedException("invalid access token"));
+            return new Response(
+                    accessToken.getValue(),
+                    determineExpiration(getExpiration(accessToken), claims.getExpirationTime()),
+                    buildRefreshToken(refreshToken, existingRefreshToken));
+        }
+    }
+
+    private class OAuth2WithOidcExtensionsCodeFlow
+            implements AuthorizationCodeFlow
+    {
+        private final IDTokenValidator idTokenValidator;
+
+        public OAuth2WithOidcExtensionsCodeFlow()
+        {
+            idTokenValidator = new IDTokenValidator(issuer, clientId, jwsKeySelector, null);
+            idTokenValidator.setMaxClockSkew((int) maxClockSkew.roundTo(SECONDS));
+        }
+
+        @Override
+        public Request createAuthorizationRequest(String state, URI callbackUri)
+        {
+            String nonce = new Nonce().getValue();
+            return new Request(
+                    new AuthenticationRequest.Builder(CODE, scope, clientId, callbackUri)
+                            .endpointURI(authUrl)
+                            .state(new State(state))
+                            .nonce(new Nonce(hashNonce(nonce)))
+                            .build()
+                            .toURI(),
+                    Optional.of(nonce));
+        }
+
+        @Override
+        public Response getOAuth2Response(String code, URI callbackUri, Optional nonce)
+                throws ChallengeFailedException
+        {
+            if (!nonce.isPresent()) {
+                throw new ChallengeFailedException("Missing nonce");
+            }
+
+            OIDCTokenResponse tokenResponse = getTokenResponse(code, callbackUri, OIDCTokenResponse::parse);
+            OIDCTokens tokens = tokenResponse.getOIDCTokens();
+            validateTokens(tokens, nonce);
+            return toResponse(tokens, Optional.empty());
+        }
+
+        @Override
+        public Response refreshTokens(String refreshToken)
+                throws ChallengeFailedException
+        {
+            OIDCTokenResponse tokenResponse = getTokenResponse(refreshToken, OIDCTokenResponse::parse);
+            OIDCTokens tokens = tokenResponse.getOIDCTokens();
+            validateTokens(tokens);
+            return toResponse(tokens, Optional.of(refreshToken));
+        }
+
+        private Response toResponse(OIDCTokens tokens, Optional existingRefreshToken)
+                throws ChallengeFailedException
+        {
+            AccessToken accessToken = tokens.getAccessToken();
+            RefreshToken refreshToken = tokens.getRefreshToken();
+            JWTClaimsSet claims = getJWTClaimsSet(accessToken.getValue()).orElseThrow(() -> new ChallengeFailedException("invalid access token"));
+            return new Response(
+                    accessToken.getValue(),
+                    determineExpiration(getExpiration(accessToken), claims.getExpirationTime()),
+                    buildRefreshToken(refreshToken, existingRefreshToken));
+        }
+
+        private void validateTokens(OIDCTokens tokens, Optional nonce)
+                throws ChallengeFailedException
+        {
+            try {
+                IDTokenClaimsSet idToken = idTokenValidator.validate(
+                        tokens.getIDToken(),
+                        nonce.map(this::hashNonce)
+                                .map(Nonce::new)
+                                .orElse(null));
+                AccessTokenHash accessTokenHash = idToken.getAccessTokenHash();
+                if (accessTokenHash != null) {
+                    AccessTokenValidator.validate(tokens.getAccessToken(), ((JWSHeader) tokens.getIDToken().getHeader()).getAlgorithm(), accessTokenHash);
+                }
+            }
+            catch (BadJOSEException | JOSEException | InvalidHashException e) {
+                throw new ChallengeFailedException("Cannot validate tokens", e);
+            }
+        }
+
+        private void validateTokens(OIDCTokens tokens)
+                throws ChallengeFailedException
+        {
+            validateTokens(tokens, Optional.empty());
+        }
+
+        private String hashNonce(String nonce)
+        {
+            return sha256()
+                    .hashString(nonce, UTF_8)
+                    .toString();
+        }
+    }
+
+    private  T getTokenResponse(String code, URI callbackUri, NimbusAirliftHttpClient.Parser parser)
+            throws ChallengeFailedException
+    {
+        return getTokenResponse(new AuthorizationCodeGrant(new AuthorizationCode(code), callbackUri), parser);
+    }
+
+    private  T getTokenResponse(String refreshToken, NimbusAirliftHttpClient.Parser parser)
+            throws ChallengeFailedException
+    {
+        return getTokenResponse(new RefreshTokenGrant(new RefreshToken(refreshToken)), parser);
+    }
+
+    private  T getTokenResponse(AuthorizationGrant authorizationGrant, NimbusAirliftHttpClient.Parser parser)
+            throws ChallengeFailedException
+    {
+        T tokenResponse = httpClient.execute(new TokenRequest(tokenUrl, clientAuth, authorizationGrant, scope), parser);
+        if (!tokenResponse.indicatesSuccess()) {
+            throw new ChallengeFailedException("Error while fetching access token: " + tokenResponse.toErrorResponse().toJSONObject());
+        }
+        return tokenResponse;
+    }
+
+    private Optional getJWTClaimsSet(String accessToken)
+    {
+        if (userinfoUrl.isPresent()) {
+            return queryUserInfo(accessToken);
+        }
+        return parseAccessToken(accessToken);
+    }
+
+    private Optional queryUserInfo(String accessToken)
+    {
+        try {
+            UserInfoResponse response = httpClient.execute(new UserInfoRequest(userinfoUrl.get(), new BearerAccessToken(accessToken)), this::parse);
+            if (!response.indicatesSuccess()) {
+                LOG.error("Received bad response from userinfo endpoint: " + response.toErrorResponse().getErrorObject());
+                return Optional.empty();
+            }
+            return Optional.of(response.toSuccessResponse().getUserInfo().toJWTClaimsSet());
+        }
+        catch (ParseException | RuntimeException e) {
+            LOG.error(e, "Received bad response from userinfo endpoint");
+            return Optional.empty();
+        }
+    }
+
+    // Using this parsing method for our /userinfo response from the IdP in order to allow for different principal
+    // fields as defined, and in the absence of the `sub` claim. This is a "hack" solution to alter the claims
+    // present in the response before calling the parser provided by the oidc sdk, which fails hard if the
+    // `sub` claim is missing. Note we also have to offload audience verification to this method since it
+    // is not handled in the library
+    public UserInfoResponse parse(HTTPResponse httpResponse)
+            throws ParseException
+    {
+        JSONObject body = httpResponse.getContentAsJSONObject();
+
+        String principal = (String) body.get(principalField);
+        if (principal == null) {
+            throw new ParseException(String.format("/userinfo response missing principal field %s", principalField));
+        }
+
+        if (!principalField.equals("sub") && body.get("sub") == null) {
+            body.put("sub", principal);
+            httpResponse.setBody(body.toJSONString());
+        }
+
+        Object audClaim = body.get("aud");
+        List audiences;
+
+        if (audClaim instanceof String) {
+            audiences = List.of((String) audClaim);
+        }
+        else if (audClaim instanceof List) {
+            audiences = ((List) audClaim).stream()
+                    .filter(String.class::isInstance)
+                    .map(String.class::cast)
+                    .collect(toImmutableList());
+        }
+        else {
+            throw new ParseException("Unsupported or missing 'aud' claim type in /userinfo response");
+        }
+
+        if (!(audiences.contains(clientId.getValue()) || !Collections.disjoint(audiences, accessTokenAudiences))) {
+            throw new ParseException("Invalid audience in /userinfo response");
+        }
+
+        return (httpResponse.getStatusCode() == 200)
+                ? UserInfoSuccessResponse.parse(httpResponse)
+                : UserInfoErrorResponse.parse(httpResponse);
+    }
+
+    private Optional parseAccessToken(String accessToken)
+    {
+        try {
+            return Optional.of(accessTokenProcessor.process(accessToken, null));
+        }
+        catch (java.text.ParseException | BadJOSEException | JOSEException e) {
+            LOG.error(e, "Failed to parse JWT access token");
+            return Optional.empty();
+        }
+    }
+
+    private static Instant determineExpiration(Optional validUntil, Date expiration)
+            throws ChallengeFailedException
+    {
+        if (validUntil.isPresent()) {
+            if (expiration != null) {
+                return Ordering.natural().min(validUntil.get(), expiration.toInstant());
+            }
+
+            return validUntil.get();
+        }
+
+        if (expiration != null) {
+            return expiration.toInstant();
+        }
+
+        throw new ChallengeFailedException("no valid expiration date");
+    }
+
+    private Optional buildRefreshToken(RefreshToken refreshToken, Optional existingRefreshToken)
+    {
+        Optional firstOption = Optional.ofNullable(refreshToken)
+                .map(RefreshToken::getValue);
+
+        if (firstOption.isPresent()) {
+            return firstOption;
+        }
+        else if (existingRefreshToken.isPresent()) {
+            return existingRefreshToken;
+        }
+        else {
+            return Optional.empty();
+        }
+    }
+
+    private static Optional getExpiration(AccessToken accessToken)
+    {
+        return accessToken.getLifetime() != 0 ? Optional.of(Instant.now().plusSeconds(accessToken.getLifetime())) : Optional.empty();
+    }
+}
diff --git a/presto-main/src/main/java/com/facebook/presto/server/security/oauth2/NonceCookie.java b/presto-main/src/main/java/com/facebook/presto/server/security/oauth2/NonceCookie.java
new file mode 100644
index 0000000000000..d80e2cacc6941
--- /dev/null
+++ b/presto-main/src/main/java/com/facebook/presto/server/security/oauth2/NonceCookie.java
@@ -0,0 +1,88 @@
+/*
+ * Licensed 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 com.facebook.presto.server.security.oauth2;
+
+import jakarta.ws.rs.core.Cookie;
+import jakarta.ws.rs.core.NewCookie;
+import org.apache.commons.lang3.StringUtils;
+
+import java.time.Instant;
+import java.util.Date;
+import java.util.Optional;
+
+import static com.facebook.presto.server.security.oauth2.OAuth2CallbackResource.CALLBACK_ENDPOINT;
+import static com.google.common.base.Predicates.not;
+import static jakarta.ws.rs.core.Cookie.DEFAULT_VERSION;
+import static jakarta.ws.rs.core.NewCookie.DEFAULT_MAX_AGE;
+
+public final class NonceCookie
+{
+    // prefix according to: https://tools.ietf.org/html/draft-ietf-httpbis-rfc6265bis-05#section-4.1.3.1
+    public static final String NONCE_COOKIE = "__Secure-Presto-Nonce";
+
+    private NonceCookie() {}
+
+    public static NewCookie create(String nonce, Instant tokenExpiration)
+    {
+        return new NewCookie(
+                NONCE_COOKIE,
+                nonce,
+                CALLBACK_ENDPOINT,
+                null,
+                DEFAULT_VERSION,
+                null,
+                DEFAULT_MAX_AGE,
+                Date.from(tokenExpiration),
+                true,
+                true);
+    }
+
+    public static jakarta.servlet.http.Cookie createServletCookie(String nonce, Instant tokenExpiration)
+    {
+        return toServletCookie(create(nonce, tokenExpiration));
+    }
+
+    public static jakarta.servlet.http.Cookie toServletCookie(NewCookie cookie)
+    {
+        jakarta.servlet.http.Cookie servletCookie = new jakarta.servlet.http.Cookie(cookie.getName(), cookie.getValue());
+        servletCookie.setPath(cookie.getPath());
+        servletCookie.setMaxAge(cookie.getMaxAge());
+        servletCookie.setSecure(cookie.isSecure());
+        servletCookie.setHttpOnly(cookie.isHttpOnly());
+
+        return servletCookie;
+    }
+
+    public static Optional read(Cookie cookie)
+    {
+        return Optional.ofNullable(cookie)
+                .map(Cookie::getValue)
+                .filter(not(StringUtils::isBlank));
+    }
+
+    public static NewCookie delete()
+    {
+        return new NewCookie(
+                NONCE_COOKIE,
+                "delete",
+                CALLBACK_ENDPOINT,
+                null,
+                DEFAULT_VERSION,
+                null,
+                0,
+                null,
+                true,
+                true);
+    }
+}
diff --git a/presto-main/src/main/java/com/facebook/presto/server/security/oauth2/OAuth2AuthenticationSupportModule.java b/presto-main/src/main/java/com/facebook/presto/server/security/oauth2/OAuth2AuthenticationSupportModule.java
new file mode 100644
index 0000000000000..88cd2b5e0f23d
--- /dev/null
+++ b/presto-main/src/main/java/com/facebook/presto/server/security/oauth2/OAuth2AuthenticationSupportModule.java
@@ -0,0 +1,33 @@
+/*
+ * Licensed 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 com.facebook.presto.server.security.oauth2;
+
+import com.facebook.airlift.configuration.AbstractConfigurationAwareModule;
+import com.google.inject.Binder;
+import com.google.inject.Scopes;
+
+import static com.facebook.airlift.jaxrs.JaxrsBinder.jaxrsBinder;
+
+public class OAuth2AuthenticationSupportModule
+        extends AbstractConfigurationAwareModule
+{
+    @Override
+    protected void setup(Binder binder)
+    {
+        binder.bind(OAuth2TokenExchange.class).in(Scopes.SINGLETON);
+        binder.bind(OAuth2TokenHandler.class).to(OAuth2TokenExchange.class).in(Scopes.SINGLETON);
+        jaxrsBinder(binder).bind(OAuth2TokenExchangeResource.class);
+        install(new OAuth2ServiceModule());
+    }
+}
diff --git a/presto-main/src/main/java/com/facebook/presto/server/security/oauth2/OAuth2Authenticator.java b/presto-main/src/main/java/com/facebook/presto/server/security/oauth2/OAuth2Authenticator.java
new file mode 100644
index 0000000000000..83c2de302c287
--- /dev/null
+++ b/presto-main/src/main/java/com/facebook/presto/server/security/oauth2/OAuth2Authenticator.java
@@ -0,0 +1,146 @@
+/*
+ * Licensed 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 com.facebook.presto.server.security.oauth2;
+
+import com.facebook.airlift.http.server.AuthenticationException;
+import com.facebook.airlift.http.server.Authenticator;
+import com.facebook.airlift.http.server.BasicPrincipal;
+import com.facebook.airlift.log.Logger;
+import jakarta.servlet.http.Cookie;
+import jakarta.servlet.http.HttpServletRequest;
+import org.apache.commons.lang3.StringUtils;
+
+import javax.inject.Inject;
+
+import java.net.URI;
+import java.security.Principal;
+import java.sql.Date;
+import java.time.Instant;
+import java.util.Arrays;
+import java.util.Map;
+import java.util.Optional;
+import java.util.UUID;
+
+import static com.facebook.presto.server.security.oauth2.OAuth2TokenExchangeResource.getInitiateUri;
+import static com.facebook.presto.server.security.oauth2.OAuth2TokenExchangeResource.getTokenUri;
+import static com.facebook.presto.server.security.oauth2.OAuth2Utils.getSchemeUriBuilder;
+import static com.facebook.presto.server.security.oauth2.OAuthWebUiCookie.OAUTH2_COOKIE;
+import static com.google.common.base.Strings.nullToEmpty;
+import static com.google.common.net.HttpHeaders.AUTHORIZATION;
+import static java.lang.String.format;
+import static java.util.Objects.requireNonNull;
+
+public class OAuth2Authenticator
+        implements Authenticator
+{
+    private static final Logger Log = Logger.get(OAuth2Authenticator.class);
+    private final String principalField;
+    private final OAuth2Client client;
+    private final TokenPairSerializer tokenPairSerializer;
+    private final TokenRefresher tokenRefresher;
+
+    @Inject
+    public OAuth2Authenticator(OAuth2Client client, OAuth2Config config, TokenRefresher tokenRefresher, TokenPairSerializer tokenPairSerializer)
+    {
+        this.client = requireNonNull(client, "service is null");
+        this.principalField = config.getPrincipalField();
+        requireNonNull(config, "oauth2Config is null");
+        this.tokenRefresher = requireNonNull(tokenRefresher, "tokenRefresher is null");
+        this.tokenPairSerializer = requireNonNull(tokenPairSerializer, "tokenPairSerializer is null");
+    }
+
+    public Principal authenticate(HttpServletRequest request) throws AuthenticationException
+    {
+        String token = extractToken(request);
+        TokenPairSerializer.TokenPair tokenPair;
+        try {
+            tokenPair = tokenPairSerializer.deserialize(token);
+        }
+        catch (IllegalArgumentException e) {
+            Log.error(e, "Failed to deserialize the OAuth token");
+            throw needAuthentication(request, Optional.empty(), "Invalid Credentials");
+        }
+
+        if (tokenPair.getExpiration().before(Date.from(Instant.now()))) {
+            throw needAuthentication(request, Optional.of(token), "Invalid Credentials");
+        }
+        Optional> claims = client.getClaims(tokenPair.getAccessToken());
+
+        if (!claims.isPresent()) {
+            throw needAuthentication(request, Optional.ofNullable(token), "Invalid Credentials");
+        }
+        String principal = (String) claims.get().get(principalField);
+        if (StringUtils.isEmpty(principal)) {
+            Log.warn("The subject is not present we need to authenticate");
+            needAuthentication(request, Optional.empty(), "Invalid Credentials");
+        }
+
+        return new BasicPrincipal(principal);
+    }
+
+    public String extractToken(HttpServletRequest request) throws AuthenticationException
+    {
+        Optional cookieToken = this.extractTokenFromCookie(request);
+        Optional headerToken = this.extractTokenFromHeader(request);
+
+        if (!cookieToken.isPresent() && !headerToken.isPresent()) {
+            throw needAuthentication(request, Optional.empty(), "Invalid Credentials");
+        }
+
+        return cookieToken.orElseGet(() -> headerToken.get());
+    }
+
+    public Optional extractTokenFromHeader(HttpServletRequest request)
+    {
+        String authHeader = nullToEmpty(request.getHeader(AUTHORIZATION));
+        int space = authHeader.indexOf(' ');
+        if ((space < 0) || !authHeader.substring(0, space).equalsIgnoreCase("bearer")) {
+            return Optional.empty();
+        }
+
+        return Optional.ofNullable(authHeader.substring(space + 1).trim())
+                .filter(t -> !t.isEmpty());
+    }
+
+    public static Optional extractTokenFromCookie(HttpServletRequest request)
+    {
+        Cookie[] cookies = Optional.ofNullable(request.getCookies()).orElse(new Cookie[0]);
+        return Optional.ofNullable(Arrays.stream(cookies)
+                .filter(cookie -> cookie.getName().equals(OAUTH2_COOKIE))
+                .findFirst()
+                .map(c -> c.getValue())
+                .orElse(null));
+    }
+
+    private AuthenticationException needAuthentication(HttpServletRequest request, Optional currentToken, String message)
+    {
+        URI baseUri = getSchemeUriBuilder(request).build();
+        return currentToken
+                .map(tokenPairSerializer::deserialize)
+                .flatMap(tokenRefresher::refreshToken)
+                .map(refreshId -> baseUri.resolve(getTokenUri(refreshId)))
+                .map(tokenUri -> new AuthenticationException(message, format("Bearer x_token_server=\"%s\"", tokenUri)))
+                .orElseGet(() -> buildNeedAuthentication(request, message));
+    }
+
+    private AuthenticationException buildNeedAuthentication(HttpServletRequest request, String message)
+    {
+        UUID authId = UUID.randomUUID();
+        URI baseUri = getSchemeUriBuilder(request).build();
+        URI initiateUri = baseUri.resolve(getInitiateUri(authId));
+        URI tokenUri = baseUri.resolve(getTokenUri(authId));
+
+        return new AuthenticationException(message, format("Bearer x_redirect_server=\"%s\", x_token_server=\"%s\"", initiateUri, tokenUri));
+    }
+}
diff --git a/presto-main/src/main/java/com/facebook/presto/server/security/oauth2/OAuth2CallbackResource.java b/presto-main/src/main/java/com/facebook/presto/server/security/oauth2/OAuth2CallbackResource.java
new file mode 100644
index 0000000000000..42d52450837d2
--- /dev/null
+++ b/presto-main/src/main/java/com/facebook/presto/server/security/oauth2/OAuth2CallbackResource.java
@@ -0,0 +1,80 @@
+/*
+ * Licensed 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 com.facebook.presto.server.security.oauth2;
+
+import com.facebook.airlift.log.Logger;
+import jakarta.servlet.http.HttpServletRequest;
+import jakarta.ws.rs.CookieParam;
+import jakarta.ws.rs.GET;
+import jakarta.ws.rs.Path;
+import jakarta.ws.rs.Produces;
+import jakarta.ws.rs.QueryParam;
+import jakarta.ws.rs.core.Context;
+import jakarta.ws.rs.core.Cookie;
+import jakarta.ws.rs.core.Response;
+import jakarta.ws.rs.core.UriBuilder;
+
+import javax.inject.Inject;
+
+import static com.facebook.presto.server.security.oauth2.NonceCookie.NONCE_COOKIE;
+import static com.facebook.presto.server.security.oauth2.OAuth2Utils.getSchemeUriBuilder;
+import static jakarta.ws.rs.core.MediaType.TEXT_HTML;
+import static jakarta.ws.rs.core.Response.Status.BAD_REQUEST;
+import static java.util.Objects.requireNonNull;
+
+@Path(OAuth2CallbackResource.CALLBACK_ENDPOINT)
+public class OAuth2CallbackResource
+{
+    private static final Logger LOG = Logger.get(OAuth2CallbackResource.class);
+
+    public static final String CALLBACK_ENDPOINT = "/oauth2/callback";
+
+    private final OAuth2Service service;
+
+    @Inject
+    public OAuth2CallbackResource(OAuth2Service service)
+    {
+        this.service = requireNonNull(service, "service is null");
+    }
+
+    @GET
+    @Produces(TEXT_HTML)
+    public Response callback(
+            @QueryParam("state") String state,
+            @QueryParam("code") String code,
+            @QueryParam("error") String error,
+            @QueryParam("error_description") String errorDescription,
+            @QueryParam("error_uri") String errorUri,
+            @CookieParam(NONCE_COOKIE) Cookie nonce,
+            @Context HttpServletRequest request)
+    {
+        if (error != null) {
+            return service.handleOAuth2Error(state, error, errorDescription, errorUri);
+        }
+
+        try {
+            requireNonNull(state, "state is null");
+            requireNonNull(code, "code is null");
+            UriBuilder builder = getSchemeUriBuilder(request);
+            return service.finishOAuth2Challenge(state, code, builder.build().resolve(CALLBACK_ENDPOINT), NonceCookie.read(nonce), request);
+        }
+        catch (RuntimeException e) {
+            LOG.error(e, "Authentication response could not be verified: state=%s", state);
+            return Response.status(BAD_REQUEST)
+                    .cookie(NonceCookie.delete())
+                    .entity(service.getInternalFailureHtml("Authentication response could not be verified"))
+                    .build();
+        }
+    }
+}
diff --git a/presto-main/src/main/java/com/facebook/presto/server/security/oauth2/OAuth2Client.java b/presto-main/src/main/java/com/facebook/presto/server/security/oauth2/OAuth2Client.java
new file mode 100644
index 0000000000000..fd68a2a5a2e06
--- /dev/null
+++ b/presto-main/src/main/java/com/facebook/presto/server/security/oauth2/OAuth2Client.java
@@ -0,0 +1,88 @@
+/*
+ * Licensed 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 com.facebook.presto.server.security.oauth2;
+
+import java.net.URI;
+import java.time.Instant;
+import java.util.Map;
+import java.util.Optional;
+
+import static java.util.Objects.requireNonNull;
+
+public interface OAuth2Client
+{
+    void load();
+
+    Request createAuthorizationRequest(String state, URI callbackUri);
+
+    Response getOAuth2Response(String code, URI callbackUri, Optional nonce)
+            throws ChallengeFailedException;
+
+    Optional> getClaims(String accessToken);
+
+    Response refreshTokens(String refreshToken)
+            throws ChallengeFailedException;
+
+    class Request
+    {
+        private final URI authorizationUri;
+        private final Optional nonce;
+
+        public Request(URI authorizationUri, Optional nonce)
+        {
+            this.authorizationUri = requireNonNull(authorizationUri, "authorizationUri is null");
+            this.nonce = requireNonNull(nonce, "nonce is null");
+        }
+
+        public URI getAuthorizationUri()
+        {
+            return authorizationUri;
+        }
+
+        public Optional getNonce()
+        {
+            return nonce;
+        }
+    }
+
+    class Response
+    {
+        private final String accessToken;
+        private final Instant expiration;
+
+        private final Optional refreshToken;
+
+        public Response(String accessToken, Instant expiration, Optional refreshToken)
+        {
+            this.accessToken = requireNonNull(accessToken, "accessToken is null");
+            this.expiration = requireNonNull(expiration, "expiration is null");
+            this.refreshToken = requireNonNull(refreshToken, "refreshToken is null");
+        }
+
+        public String getAccessToken()
+        {
+            return accessToken;
+        }
+
+        public Instant getExpiration()
+        {
+            return expiration;
+        }
+
+        public Optional getRefreshToken()
+        {
+            return refreshToken;
+        }
+    }
+}
diff --git a/presto-main/src/main/java/com/facebook/presto/server/security/oauth2/OAuth2Config.java b/presto-main/src/main/java/com/facebook/presto/server/security/oauth2/OAuth2Config.java
new file mode 100644
index 0000000000000..b1b0f9513b2f7
--- /dev/null
+++ b/presto-main/src/main/java/com/facebook/presto/server/security/oauth2/OAuth2Config.java
@@ -0,0 +1,253 @@
+/*
+ * Licensed 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 com.facebook.presto.server.security.oauth2;
+
+import com.facebook.airlift.configuration.Config;
+import com.facebook.airlift.configuration.ConfigDescription;
+import com.facebook.airlift.configuration.ConfigSecuritySensitive;
+import com.facebook.airlift.units.Duration;
+import com.facebook.airlift.units.MinDuration;
+import com.google.common.base.Splitter;
+import com.google.common.collect.ImmutableList;
+import com.google.common.collect.ImmutableSet;
+import jakarta.validation.constraints.NotNull;
+
+import java.io.File;
+import java.util.Collections;
+import java.util.List;
+import java.util.Optional;
+import java.util.Set;
+import java.util.concurrent.TimeUnit;
+
+import static com.facebook.presto.server.security.oauth2.OAuth2Service.OPENID_SCOPE;
+import static com.google.common.base.Strings.nullToEmpty;
+
+public class OAuth2Config
+{
+    private Optional stateKey = Optional.empty();
+    private String issuer;
+    private String clientId;
+    private String clientSecret;
+    private Set scopes = ImmutableSet.of(OPENID_SCOPE);
+    private String principalField = "sub";
+    private Optional groupsField = Optional.empty();
+    private List additionalAudiences = Collections.emptyList();
+    private Duration challengeTimeout = new Duration(15, TimeUnit.MINUTES);
+    private Duration maxClockSkew = new Duration(1, TimeUnit.MINUTES);
+    private Optional userMappingPattern = Optional.empty();
+    private Optional userMappingFile = Optional.empty();
+    private boolean enableRefreshTokens;
+    private boolean enableDiscovery = true;
+
+    public Optional getStateKey()
+    {
+        return stateKey;
+    }
+
+    @Config("http-server.authentication.oauth2.state-key")
+    @ConfigDescription("A secret key used by HMAC algorithm to sign the state parameter")
+    public OAuth2Config setStateKey(String stateKey)
+    {
+        this.stateKey = Optional.ofNullable(stateKey);
+        return this;
+    }
+
+    @NotNull
+    public String getIssuer()
+    {
+        return issuer;
+    }
+
+    @Config("http-server.authentication.oauth2.issuer")
+    @ConfigDescription("The required issuer of a token")
+    public OAuth2Config setIssuer(String issuer)
+    {
+        this.issuer = issuer;
+        return this;
+    }
+
+    @NotNull
+    public String getClientId()
+    {
+        return clientId;
+    }
+
+    @Config("http-server.authentication.oauth2.client-id")
+    @ConfigDescription("Client ID")
+    public OAuth2Config setClientId(String clientId)
+    {
+        this.clientId = clientId;
+        return this;
+    }
+
+    @NotNull
+    public String getClientSecret()
+    {
+        return clientSecret;
+    }
+
+    @Config("http-server.authentication.oauth2.client-secret")
+    @ConfigSecuritySensitive
+    @ConfigDescription("Client secret")
+    public OAuth2Config setClientSecret(String clientSecret)
+    {
+        this.clientSecret = clientSecret;
+        return this;
+    }
+
+    @NotNull
+    public List getAdditionalAudiences()
+    {
+        return additionalAudiences;
+    }
+
+    public OAuth2Config setAdditionalAudiences(List additionalAudiences)
+    {
+        this.additionalAudiences = ImmutableList.copyOf(additionalAudiences);
+        return this;
+    }
+
+    @Config("http-server.authentication.oauth2.additional-audiences")
+    @ConfigDescription("Additional audiences to trust in addition to the Client ID")
+    public OAuth2Config setAdditionalAudiences(String additionalAudiences)
+    {
+        Splitter splitter = Splitter.on(',').trimResults().omitEmptyStrings();
+        this.additionalAudiences = ImmutableList.copyOf(splitter.split(nullToEmpty(additionalAudiences)));
+        return this;
+    }
+
+    @NotNull
+    public Set getScopes()
+    {
+        return scopes;
+    }
+
+    @Config("http-server.authentication.oauth2.scopes")
+    @ConfigDescription("Scopes requested by the server during OAuth2 authorization challenge")
+    public OAuth2Config setScopes(String scopes)
+    {
+        Splitter splitter = Splitter.on(',').trimResults().omitEmptyStrings();
+        this.scopes = ImmutableSet.copyOf(splitter.split(nullToEmpty(scopes)));
+        return this;
+    }
+
+    @NotNull
+    public String getPrincipalField()
+    {
+        return principalField;
+    }
+
+    @Config("http-server.authentication.oauth2.principal-field")
+    @ConfigDescription("The claim to use as the principal")
+    public OAuth2Config setPrincipalField(String principalField)
+    {
+        this.principalField = principalField;
+        return this;
+    }
+
+    public Optional getGroupsField()
+    {
+        return groupsField;
+    }
+
+    @Config("http-server.authentication.oauth2.groups-field")
+    @ConfigDescription("Groups field in the claim")
+    public OAuth2Config setGroupsField(String groupsField)
+    {
+        this.groupsField = Optional.ofNullable(groupsField);
+        return this;
+    }
+
+    @MinDuration("1ms")
+    @NotNull
+    public Duration getChallengeTimeout()
+    {
+        return challengeTimeout;
+    }
+
+    @Config("http-server.authentication.oauth2.challenge-timeout")
+    @ConfigDescription("Maximum duration of OAuth2 authorization challenge")
+    public OAuth2Config setChallengeTimeout(Duration challengeTimeout)
+    {
+        this.challengeTimeout = challengeTimeout;
+        return this;
+    }
+
+    @MinDuration("0s")
+    @NotNull
+    public Duration getMaxClockSkew()
+    {
+        return maxClockSkew;
+    }
+
+    @Config("http-server.authentication.oauth2.max-clock-skew")
+    @ConfigDescription("Max clock skew between the Authorization Server and the coordinator")
+    public OAuth2Config setMaxClockSkew(Duration maxClockSkew)
+    {
+        this.maxClockSkew = maxClockSkew;
+        return this;
+    }
+
+    public Optional getUserMappingPattern()
+    {
+        return userMappingPattern;
+    }
+
+    @Config("http-server.authentication.oauth2.user-mapping.pattern")
+    @ConfigDescription("Regex to match against user name")
+    public OAuth2Config setUserMappingPattern(String userMappingPattern)
+    {
+        this.userMappingPattern = Optional.ofNullable(userMappingPattern);
+        return this;
+    }
+
+    public Optional getUserMappingFile()
+    {
+        return userMappingFile;
+    }
+
+    @Config("http-server.authentication.oauth2.user-mapping.file")
+    @ConfigDescription("File containing rules for mapping user")
+    public OAuth2Config setUserMappingFile(File userMappingFile)
+    {
+        this.userMappingFile = Optional.ofNullable(userMappingFile);
+        return this;
+    }
+
+    public boolean isEnableRefreshTokens()
+    {
+        return enableRefreshTokens;
+    }
+
+    @Config("http-server.authentication.oauth2.refresh-tokens")
+    @ConfigDescription("Enables OpenID refresh tokens usage")
+    public OAuth2Config setEnableRefreshTokens(boolean enableRefreshTokens)
+    {
+        this.enableRefreshTokens = enableRefreshTokens;
+        return this;
+    }
+
+    public boolean isEnableDiscovery()
+    {
+        return enableDiscovery;
+    }
+
+    @Config("http-server.authentication.oauth2.oidc.discovery")
+    @ConfigDescription("Enable OpenID Provider Issuer discovery")
+    public OAuth2Config setEnableDiscovery(boolean enableDiscovery)
+    {
+        this.enableDiscovery = enableDiscovery;
+        return this;
+    }
+}
diff --git a/presto-main/src/main/java/com/facebook/presto/server/security/oauth2/OAuth2ErrorCode.java b/presto-main/src/main/java/com/facebook/presto/server/security/oauth2/OAuth2ErrorCode.java
new file mode 100644
index 0000000000000..058ce8be92b1e
--- /dev/null
+++ b/presto-main/src/main/java/com/facebook/presto/server/security/oauth2/OAuth2ErrorCode.java
@@ -0,0 +1,46 @@
+/*
+ * Licensed 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 com.facebook.presto.server.security.oauth2;
+
+import java.util.Arrays;
+
+public enum OAuth2ErrorCode
+{
+    ACCESS_DENIED("access_denied", "OAuth2 server denied the login"),
+    UNAUTHORIZED_CLIENT("unauthorized_client", "OAuth2 server does not allow request from this Presto server"),
+    SERVER_ERROR("server_error", "OAuth2 server had a failure"),
+    TEMPORARILY_UNAVAILABLE("temporarily_unavailable", "OAuth2 server is temporarily unavailable");
+
+    private final String code;
+    private final String message;
+
+    OAuth2ErrorCode(String code, String message)
+    {
+        this.code = code;
+        this.message = message;
+    }
+
+    public static OAuth2ErrorCode fromString(String codeStr)
+    {
+        return Arrays.stream(OAuth2ErrorCode.values())
+                .filter(value -> codeStr.equalsIgnoreCase(value.code))
+                .findFirst()
+                .orElseThrow(() -> new IllegalArgumentException("No enum constant " + OAuth2ErrorCode.class.getCanonicalName() + "." + codeStr));
+    }
+
+    public String getMessage()
+    {
+        return this.message;
+    }
+}
diff --git a/presto-main/src/main/java/com/facebook/presto/server/security/oauth2/OAuth2ServerConfigProvider.java b/presto-main/src/main/java/com/facebook/presto/server/security/oauth2/OAuth2ServerConfigProvider.java
new file mode 100644
index 0000000000000..9455e5622605d
--- /dev/null
+++ b/presto-main/src/main/java/com/facebook/presto/server/security/oauth2/OAuth2ServerConfigProvider.java
@@ -0,0 +1,67 @@
+/*
+ * Licensed 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 com.facebook.presto.server.security.oauth2;
+
+import java.net.URI;
+import java.util.Optional;
+
+import static java.util.Objects.requireNonNull;
+
+public interface OAuth2ServerConfigProvider
+{
+    OAuth2ServerConfig get();
+
+    class OAuth2ServerConfig
+    {
+        private final Optional accessTokenIssuer;
+        private final URI authUrl;
+        private final URI tokenUrl;
+        private final URI jwksUrl;
+        private final Optional userinfoUrl;
+
+        public OAuth2ServerConfig(Optional accessTokenIssuer, URI authUrl, URI tokenUrl, URI jwksUrl, Optional userinfoUrl)
+        {
+            this.accessTokenIssuer = requireNonNull(accessTokenIssuer, "accessTokenIssuer is null");
+            this.authUrl = requireNonNull(authUrl, "authUrl is null");
+            this.tokenUrl = requireNonNull(tokenUrl, "tokenUrl is null");
+            this.jwksUrl = requireNonNull(jwksUrl, "jwksUrl is null");
+            this.userinfoUrl = requireNonNull(userinfoUrl, "userinfoUrl is null");
+        }
+
+        public Optional getAccessTokenIssuer()
+        {
+            return accessTokenIssuer;
+        }
+
+        public URI getAuthUrl()
+        {
+            return authUrl;
+        }
+
+        public URI getTokenUrl()
+        {
+            return tokenUrl;
+        }
+
+        public URI getJwksUrl()
+        {
+            return jwksUrl;
+        }
+
+        public Optional getUserinfoUrl()
+        {
+            return userinfoUrl;
+        }
+    }
+}
diff --git a/presto-main/src/main/java/com/facebook/presto/server/security/oauth2/OAuth2Service.java b/presto-main/src/main/java/com/facebook/presto/server/security/oauth2/OAuth2Service.java
new file mode 100644
index 0000000000000..78616a7d8cad5
--- /dev/null
+++ b/presto-main/src/main/java/com/facebook/presto/server/security/oauth2/OAuth2Service.java
@@ -0,0 +1,285 @@
+/*
+ * Licensed 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 com.facebook.presto.server.security.oauth2;
+
+import com.facebook.airlift.log.Logger;
+import com.google.common.io.Resources;
+import io.jsonwebtoken.Claims;
+import io.jsonwebtoken.JwtParser;
+import jakarta.servlet.http.HttpServletRequest;
+import jakarta.servlet.http.HttpServletResponse;
+import jakarta.ws.rs.core.HttpHeaders;
+import jakarta.ws.rs.core.Response;
+import jakarta.ws.rs.core.UriBuilder;
+
+import javax.inject.Inject;
+
+import java.io.IOException;
+import java.net.URI;
+import java.security.Key;
+import java.security.SecureRandom;
+import java.time.Duration;
+import java.time.Instant;
+import java.time.temporal.TemporalAmount;
+import java.util.Date;
+import java.util.Optional;
+import java.util.Random;
+
+import static com.facebook.presto.server.security.oauth2.JwtUtil.newJwtBuilder;
+import static com.facebook.presto.server.security.oauth2.JwtUtil.newJwtParserBuilder;
+import static com.facebook.presto.server.security.oauth2.OAuth2Utils.getSchemeUriBuilder;
+import static com.facebook.presto.server.security.oauth2.TokenPairSerializer.TokenPair.fromOAuth2Response;
+import static com.google.common.base.Strings.nullToEmpty;
+import static com.google.common.base.Verify.verify;
+import static com.google.common.hash.Hashing.sha256;
+import static io.jsonwebtoken.security.Keys.hmacShaKeyFor;
+import static jakarta.ws.rs.core.Response.Status.BAD_REQUEST;
+import static jakarta.ws.rs.core.Response.Status.FORBIDDEN;
+import static java.lang.String.format;
+import static java.nio.charset.StandardCharsets.UTF_8;
+import static java.time.Instant.now;
+import static java.util.Objects.requireNonNull;
+
+public class OAuth2Service
+{
+    private static final Logger logger = Logger.get(OAuth2Service.class);
+
+    public static final String OPENID_SCOPE = "openid";
+
+    private static final String STATE_AUDIENCE_UI = "presto_oauth_ui";
+    private static final String FAILURE_REPLACEMENT_TEXT = "";
+    private static final Random SECURE_RANDOM = new SecureRandom();
+    public static final String HANDLER_STATE_CLAIM = "handler_state";
+
+    private final OAuth2Client client;
+    private final Optional tokenExpiration;
+    private final TokenPairSerializer tokenPairSerializer;
+
+    private final String successHtml;
+    private final String failureHtml;
+
+    private final TemporalAmount challengeTimeout;
+    private final Key stateHmac;
+    private final JwtParser jwtParser;
+
+    private final OAuth2TokenHandler tokenHandler;
+
+    @Inject
+    public OAuth2Service(
+            OAuth2Client client,
+            OAuth2Config oauth2Config,
+            OAuth2TokenHandler tokenHandler,
+            TokenPairSerializer tokenPairSerializer,
+            @ForRefreshTokens Optional tokenExpiration)
+            throws IOException
+    {
+        this.client = requireNonNull(client, "client is null");
+        requireNonNull(oauth2Config, "oauth2Config is null");
+
+        this.successHtml = Resources.toString(Resources.getResource(getClass(), "/webapp/oauth2/success.html"), UTF_8);
+        this.failureHtml = Resources.toString(Resources.getResource(getClass(), "/webapp/oauth2/failure.html"), UTF_8);
+        verify(failureHtml.contains(FAILURE_REPLACEMENT_TEXT), "failure.html does not contain the replacement text");
+
+        this.challengeTimeout = Duration.ofMillis(oauth2Config.getChallengeTimeout().toMillis());
+        this.stateHmac = hmacShaKeyFor(oauth2Config.getStateKey()
+                .map(key -> sha256().hashString(key, UTF_8).asBytes())
+                .orElseGet(() -> secureRandomBytes(32)));
+        this.jwtParser = newJwtParserBuilder()
+                .setSigningKey(stateHmac)
+                .requireAudience(STATE_AUDIENCE_UI)
+                .build();
+
+        this.tokenHandler = requireNonNull(tokenHandler, "tokenHandler is null");
+        this.tokenPairSerializer = requireNonNull(tokenPairSerializer, "tokenPairSerializer is null");
+
+        this.tokenExpiration = requireNonNull(tokenExpiration, "tokenExpiration is null");
+    }
+
+    public Response startOAuth2Challenge(URI callbackUri, Optional handlerState)
+    {
+        Instant challengeExpiration = now().plus(challengeTimeout);
+        String state = newJwtBuilder()
+                .signWith(stateHmac)
+                .setAudience(STATE_AUDIENCE_UI)
+                .claim(HANDLER_STATE_CLAIM, handlerState.orElse(null))
+                .setExpiration(Date.from(challengeExpiration))
+                .compact();
+
+        OAuth2Client.Request request = client.createAuthorizationRequest(state, callbackUri);
+        Response.ResponseBuilder response = Response.seeOther(request.getAuthorizationUri());
+        request.getNonce().ifPresent(nce -> response.cookie(NonceCookie.create(nce, challengeExpiration)));
+        return response.build();
+    }
+
+    public void startOAuth2Challenge(URI callbackUri, Optional handlerState, HttpServletResponse servletResponse)
+            throws IOException
+    {
+        Instant challengeExpiration = now().plus(challengeTimeout);
+
+        OAuth2Client.Request challengeRequest = this.startChallenge(callbackUri, handlerState);
+        challengeRequest.getNonce().ifPresent(nce -> servletResponse.addCookie(NonceCookie.createServletCookie(nce, challengeExpiration)));
+        servletResponseSeeOther(challengeRequest.getAuthorizationUri().toString(), servletResponse);
+    }
+
+    public void servletResponseSeeOther(String location, HttpServletResponse servletResponse)
+            throws IOException
+    {
+        // 303 is preferred over a 302 when this response is received by a POST/PUT/DELETE and the redirect should be done via a GET instead of original method
+        servletResponse.addHeader(HttpHeaders.LOCATION, location);
+        servletResponse.sendError(HttpServletResponse.SC_SEE_OTHER);
+    }
+
+    private OAuth2Client.Request startChallenge(URI callbackUri, Optional handlerState)
+    {
+        Instant challengeExpiration = now().plus(challengeTimeout);
+        String state = newJwtBuilder()
+                .signWith(stateHmac)
+                .setAudience(STATE_AUDIENCE_UI)
+                .claim(HANDLER_STATE_CLAIM, handlerState.orElse(null))
+                .setExpiration(Date.from(challengeExpiration))
+                .compact();
+
+        return client.createAuthorizationRequest(state, callbackUri);
+    }
+
+    public Response handleOAuth2Error(String state, String error, String errorDescription, String errorUri)
+    {
+        try {
+            Claims stateClaims = parseState(state);
+            Optional.ofNullable(stateClaims.get(HANDLER_STATE_CLAIM, String.class))
+                    .ifPresent(value ->
+                            tokenHandler.setTokenExchangeError(value,
+                                    format("Authentication response could not be verified: error=%s, errorDescription=%s, errorUri=%s",
+                                            error, errorDescription, errorDescription)));
+        }
+        catch (ChallengeFailedException | RuntimeException e) {
+            logger.error(e, "Authentication response could not be verified invalid state: state=%s", state);
+            return Response.status(FORBIDDEN)
+                    .entity(getInternalFailureHtml("Authentication response could not be verified"))
+                    .cookie(NonceCookie.delete())
+                    .build();
+        }
+
+        logger.error("OAuth server returned an error: error=%s, error_description=%s, error_uri=%s, state=%s", error, errorDescription, errorUri, state);
+        return Response.ok()
+                .entity(getCallbackErrorHtml(error))
+                .cookie(NonceCookie.delete())
+                .build();
+    }
+
+    public Response finishOAuth2Challenge(String state, String code, URI callbackUri, Optional nonce, HttpServletRequest request)
+    {
+        Optional handlerState;
+        try {
+            Claims stateClaims = parseState(state);
+            handlerState = Optional.ofNullable(stateClaims.get(HANDLER_STATE_CLAIM, String.class));
+        }
+        catch (ChallengeFailedException | RuntimeException e) {
+            logger.error(e, "Authentication response could not be verified invalid state: state=%s", state);
+            return Response.status(BAD_REQUEST)
+                    .entity(getInternalFailureHtml("Authentication response could not be verified"))
+                    .cookie(NonceCookie.delete())
+                    .build();
+        }
+
+        // Note: the Web UI may be disabled, so REST requests can not redirect to a success or error page inside the Web UI
+        try {
+            // fetch access token
+            OAuth2Client.Response oauth2Response = client.getOAuth2Response(code, callbackUri, nonce);
+
+            if (!handlerState.isPresent()) {
+                UriBuilder uriBuilder = getSchemeUriBuilder(request);
+                return Response
+                        .seeOther(uriBuilder.build().resolve("/ui/"))
+                        .cookie(
+                                OAuthWebUiCookie.create(
+                                        tokenPairSerializer.serialize(
+                                                fromOAuth2Response(oauth2Response)),
+                                                tokenExpiration
+                                                        .map(expiration -> Instant.now().plus(expiration))
+                                                        .orElse(oauth2Response.getExpiration())),
+                                NonceCookie.delete())
+                        .build();
+            }
+
+            tokenHandler.setAccessToken(handlerState.get(), tokenPairSerializer.serialize(fromOAuth2Response(oauth2Response)));
+
+            Response.ResponseBuilder builder = Response.ok(getSuccessHtml());
+            builder.cookie(
+                    OAuthWebUiCookie.create(
+                            tokenPairSerializer.serialize(fromOAuth2Response(oauth2Response)),
+                            tokenExpiration.map(expiration -> Instant.now().plus(expiration))
+                                    .orElse(oauth2Response.getExpiration())));
+
+            return builder.cookie(NonceCookie.delete()).build();
+        }
+        catch (ChallengeFailedException | RuntimeException e) {
+            logger.error(e, "Authentication response could not be verified: state=%s", state);
+
+            handlerState.ifPresent(value ->
+                    tokenHandler.setTokenExchangeError(value, format("Authentication response could not be verified: state=%s", value)));
+            return Response.status(BAD_REQUEST)
+                    .cookie(NonceCookie.delete())
+                    .entity(getInternalFailureHtml("Authentication response could not be verified"))
+                    .build();
+        }
+    }
+
+    private Claims parseState(String state)
+            throws ChallengeFailedException
+    {
+        try {
+            return jwtParser
+                    .parseClaimsJws(state)
+                    .getBody();
+        }
+        catch (RuntimeException e) {
+            throw new ChallengeFailedException("State validation failed", e);
+        }
+    }
+
+    public String getSuccessHtml()
+    {
+        return successHtml;
+    }
+
+    public String getCallbackErrorHtml(String errorCode)
+    {
+        return failureHtml.replace(FAILURE_REPLACEMENT_TEXT, getOAuth2ErrorMessage(errorCode));
+    }
+
+    public String getInternalFailureHtml(String errorMessage)
+    {
+        return failureHtml.replace(FAILURE_REPLACEMENT_TEXT, nullToEmpty(errorMessage));
+    }
+
+    private static byte[] secureRandomBytes(int count)
+    {
+        byte[] bytes = new byte[count];
+        SECURE_RANDOM.nextBytes(bytes);
+        return bytes;
+    }
+
+    private static String getOAuth2ErrorMessage(String errorCode)
+    {
+        try {
+            OAuth2ErrorCode code = OAuth2ErrorCode.fromString(errorCode);
+            return code.getMessage();
+        }
+        catch (IllegalArgumentException e) {
+            logger.error(e, "Unknown error code received code=%s", errorCode);
+            return "OAuth2 unknown error code: " + errorCode;
+        }
+    }
+}
diff --git a/presto-main/src/main/java/com/facebook/presto/server/security/oauth2/OAuth2ServiceModule.java b/presto-main/src/main/java/com/facebook/presto/server/security/oauth2/OAuth2ServiceModule.java
new file mode 100644
index 0000000000000..31b8740edbaf4
--- /dev/null
+++ b/presto-main/src/main/java/com/facebook/presto/server/security/oauth2/OAuth2ServiceModule.java
@@ -0,0 +1,105 @@
+/*
+ * Licensed 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 com.facebook.presto.server.security.oauth2;
+
+import com.facebook.airlift.configuration.AbstractConfigurationAwareModule;
+import com.google.inject.Binder;
+import com.google.inject.Inject;
+import com.google.inject.Key;
+import com.google.inject.Provides;
+import com.google.inject.Scopes;
+import com.google.inject.Singleton;
+
+import java.time.Duration;
+
+import static com.facebook.airlift.configuration.ConditionalModule.installModuleIf;
+import static com.facebook.airlift.configuration.ConfigBinder.configBinder;
+import static com.facebook.airlift.http.client.HttpClientBinder.httpClientBinder;
+import static com.facebook.airlift.jaxrs.JaxrsBinder.jaxrsBinder;
+import static com.facebook.presto.server.security.oauth2.TokenPairSerializer.ACCESS_TOKEN_ONLY_SERIALIZER;
+import static com.google.inject.multibindings.OptionalBinder.newOptionalBinder;
+
+public class OAuth2ServiceModule
+        extends AbstractConfigurationAwareModule
+{
+    @Override
+    protected void setup(Binder binder)
+    {
+        jaxrsBinder(binder).bind(OAuth2CallbackResource.class);
+        configBinder(binder).bindConfig(OAuth2Config.class);
+        binder.bind(OAuth2Service.class).in(Scopes.SINGLETON);
+        binder.bind(OAuth2TokenHandler.class).to(OAuth2TokenExchange.class).in(Scopes.SINGLETON);
+        binder.bind(NimbusHttpClient.class).to(NimbusAirliftHttpClient.class).in(Scopes.SINGLETON);
+        newOptionalBinder(binder, OAuth2Client.class)
+                .setDefault()
+                .to(NimbusOAuth2Client.class)
+                .in(Scopes.SINGLETON);
+        install(installModuleIf(OAuth2Config.class, OAuth2Config::isEnableDiscovery, this::bindOidcDiscovery, this::bindStaticConfiguration));
+        install(installModuleIf(OAuth2Config.class, OAuth2Config::isEnableRefreshTokens, this::enableRefreshTokens, this::disableRefreshTokens));
+        httpClientBinder(binder)
+                .bindHttpClient("oauth2-jwk", ForOAuth2.class)
+                // Reset to defaults to override InternalCommunicationModule changes to this client default configuration.
+                // Setting a keystore and/or a truststore for internal communication changes the default SSL configuration
+                // for all clients in this guice context. This does not make sense for this client which will very rarely
+                // use the same SSL configuration, so using the system default truststore makes more sense.
+                .withConfigDefaults(config -> config
+                        .setKeyStorePath(null)
+                        .setKeyStorePassword(null)
+                        .setTrustStorePath(null)
+                        .setTrustStorePassword(null));
+    }
+
+    private void enableRefreshTokens(Binder binder)
+    {
+        install(new JweTokenSerializerModule());
+    }
+
+    private void disableRefreshTokens(Binder binder)
+    {
+        binder.bind(TokenPairSerializer.class).toInstance(ACCESS_TOKEN_ONLY_SERIALIZER);
+        newOptionalBinder(binder, Key.get(Duration.class, ForRefreshTokens.class));
+    }
+
+    @Singleton
+    @Provides
+    @Inject
+    public TokenRefresher getTokenRefresher(TokenPairSerializer tokenAssembler, OAuth2TokenHandler tokenHandler, OAuth2Client oAuth2Client)
+    {
+        return new TokenRefresher(tokenAssembler, tokenHandler, oAuth2Client);
+    }
+
+    private void bindStaticConfiguration(Binder binder)
+    {
+        configBinder(binder).bindConfig(StaticOAuth2ServerConfiguration.class);
+        binder.bind(OAuth2ServerConfigProvider.class).to(StaticConfigurationProvider.class).in(Scopes.SINGLETON);
+    }
+
+    private void bindOidcDiscovery(Binder binder)
+    {
+        configBinder(binder).bindConfig(OidcDiscoveryConfig.class);
+        binder.bind(OAuth2ServerConfigProvider.class).to(OidcDiscovery.class).in(Scopes.SINGLETON);
+    }
+
+    @Override
+    public int hashCode()
+    {
+        return OAuth2ServiceModule.class.hashCode();
+    }
+
+    @Override
+    public boolean equals(Object obj)
+    {
+        return obj instanceof OAuth2ServiceModule;
+    }
+}
diff --git a/presto-main/src/main/java/com/facebook/presto/server/security/oauth2/OAuth2TokenExchange.java b/presto-main/src/main/java/com/facebook/presto/server/security/oauth2/OAuth2TokenExchange.java
new file mode 100644
index 0000000000000..95e0bdc25eb01
--- /dev/null
+++ b/presto-main/src/main/java/com/facebook/presto/server/security/oauth2/OAuth2TokenExchange.java
@@ -0,0 +1,146 @@
+/*
+ * Licensed 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 com.facebook.presto.server.security.oauth2;
+
+import com.facebook.airlift.units.Duration;
+import com.google.common.cache.CacheBuilder;
+import com.google.common.cache.CacheLoader;
+import com.google.common.cache.LoadingCache;
+import com.google.common.hash.Hashing;
+import com.google.common.util.concurrent.ListenableFuture;
+import com.google.common.util.concurrent.SettableFuture;
+import jakarta.annotation.PreDestroy;
+
+import javax.inject.Inject;
+
+import java.nio.charset.StandardCharsets;
+import java.util.Optional;
+import java.util.UUID;
+import java.util.concurrent.Future;
+import java.util.concurrent.ScheduledExecutorService;
+
+import static com.facebook.airlift.concurrent.Threads.daemonThreadsNamed;
+import static com.google.common.util.concurrent.Futures.nonCancellationPropagating;
+import static java.util.Objects.requireNonNull;
+import static java.util.concurrent.Executors.newSingleThreadScheduledExecutor;
+import static java.util.concurrent.TimeUnit.MILLISECONDS;
+import static java.util.concurrent.TimeUnit.SECONDS;
+
+public class OAuth2TokenExchange
+        implements OAuth2TokenHandler
+{
+    public static final Duration MAX_POLL_TIME = new Duration(10, SECONDS);
+    private static final TokenPoll TOKEN_POLL_TIMED_OUT = TokenPoll.error("Authentication has timed out");
+    private static final TokenPoll TOKEN_POLL_DROPPED = TokenPoll.error("Authentication has been finished by the client");
+
+    private final LoadingCache> cache;
+    private final ScheduledExecutorService executor = newSingleThreadScheduledExecutor(daemonThreadsNamed("oauth2-token-exchange"));
+
+    @Inject
+    public OAuth2TokenExchange(OAuth2Config config)
+    {
+        long challengeTimeoutMillis = config.getChallengeTimeout().toMillis();
+        this.cache = buildUnsafeCache(
+                CacheBuilder.newBuilder()
+                        .expireAfterWrite(challengeTimeoutMillis + (MAX_POLL_TIME.toMillis() * 10), MILLISECONDS)
+                        .removalListener(notification -> notification.getValue().set(TOKEN_POLL_TIMED_OUT)),
+                new CacheLoader>()
+                {
+                    @Override
+                    public SettableFuture load(String authIdHash)
+                    {
+                        SettableFuture future = SettableFuture.create();
+                        Future timeout = executor.schedule(() -> future.set(TOKEN_POLL_TIMED_OUT), challengeTimeoutMillis, MILLISECONDS);
+                        future.addListener(() -> timeout.cancel(true), executor);
+                        return future;
+                    }
+                });
+    }
+
+    private static  LoadingCache buildUnsafeCache(CacheBuilder cacheBuilder, CacheLoader cacheLoader)
+    {
+        return cacheBuilder.build(cacheLoader);
+    }
+
+    @PreDestroy
+    public void stop()
+    {
+        executor.shutdownNow();
+    }
+
+    @Override
+    public void setAccessToken(String authIdHash, String accessToken)
+    {
+        cache.getUnchecked(authIdHash).set(TokenPoll.token(accessToken));
+    }
+
+    @Override
+    public void setTokenExchangeError(String authIdHash, String message)
+    {
+        cache.getUnchecked(authIdHash).set(TokenPoll.error(message));
+    }
+
+    public ListenableFuture getTokenPoll(UUID authId)
+    {
+        return nonCancellationPropagating(cache.getUnchecked(hashAuthId(authId)));
+    }
+
+    public void dropToken(UUID authId)
+    {
+        cache.getUnchecked(hashAuthId(authId)).set(TOKEN_POLL_DROPPED);
+    }
+
+    public static String hashAuthId(UUID authId)
+    {
+        return Hashing.sha256()
+                .hashString(authId.toString(), StandardCharsets.UTF_8)
+                .toString();
+    }
+
+    public static class TokenPoll
+    {
+        private final Optional token;
+        private final Optional error;
+
+        private TokenPoll(String token, String error)
+        {
+            this.token = Optional.ofNullable(token);
+            this.error = Optional.ofNullable(error);
+        }
+
+        static TokenPoll token(String token)
+        {
+            requireNonNull(token, "token is null");
+
+            return new TokenPoll(token, null);
+        }
+
+        static TokenPoll error(String error)
+        {
+            requireNonNull(error, "error is null");
+
+            return new TokenPoll(null, error);
+        }
+
+        public Optional getToken()
+        {
+            return token;
+        }
+
+        public Optional getError()
+        {
+            return error;
+        }
+    }
+}
diff --git a/presto-main/src/main/java/com/facebook/presto/server/security/oauth2/OAuth2TokenExchangeResource.java b/presto-main/src/main/java/com/facebook/presto/server/security/oauth2/OAuth2TokenExchangeResource.java
new file mode 100644
index 0000000000000..5335712abe75e
--- /dev/null
+++ b/presto-main/src/main/java/com/facebook/presto/server/security/oauth2/OAuth2TokenExchangeResource.java
@@ -0,0 +1,141 @@
+/*
+ * Licensed 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 com.facebook.presto.server.security.oauth2;
+
+import com.facebook.airlift.json.JsonCodec;
+import com.facebook.airlift.json.JsonCodecFactory;
+import com.facebook.presto.dispatcher.DispatchExecutor;
+import com.facebook.presto.server.security.oauth2.OAuth2TokenExchange.TokenPoll;
+import com.google.common.base.VerifyException;
+import com.google.common.collect.ImmutableMap;
+import com.google.common.util.concurrent.Futures;
+import com.google.common.util.concurrent.ListenableFuture;
+import com.google.common.util.concurrent.ListeningExecutorService;
+import jakarta.servlet.http.HttpServletRequest;
+import jakarta.ws.rs.BadRequestException;
+import jakarta.ws.rs.DELETE;
+import jakarta.ws.rs.GET;
+import jakarta.ws.rs.Path;
+import jakarta.ws.rs.PathParam;
+import jakarta.ws.rs.Produces;
+import jakarta.ws.rs.container.AsyncResponse;
+import jakarta.ws.rs.container.Suspended;
+import jakarta.ws.rs.core.Context;
+import jakarta.ws.rs.core.MediaType;
+import jakarta.ws.rs.core.Response;
+import jakarta.ws.rs.core.UriBuilder;
+
+import javax.inject.Inject;
+
+import java.util.Map;
+import java.util.Optional;
+import java.util.UUID;
+
+import static com.facebook.airlift.http.server.AsyncResponseHandler.bindAsyncResponse;
+import static com.facebook.presto.server.security.oauth2.OAuth2CallbackResource.CALLBACK_ENDPOINT;
+import static com.facebook.presto.server.security.oauth2.OAuth2TokenExchange.MAX_POLL_TIME;
+import static com.facebook.presto.server.security.oauth2.OAuth2TokenExchange.hashAuthId;
+import static com.facebook.presto.server.security.oauth2.OAuth2Utils.getSchemeUriBuilder;
+import static jakarta.ws.rs.core.MediaType.APPLICATION_JSON_TYPE;
+import static java.util.Objects.requireNonNull;
+
+@Path(OAuth2TokenExchangeResource.TOKEN_ENDPOINT)
+public class OAuth2TokenExchangeResource
+{
+    public static final String TOKEN_ENDPOINT = "/oauth2/token/";
+    private static final JsonCodec> MAP_CODEC = new JsonCodecFactory().mapJsonCodec(String.class, Object.class);
+    private final OAuth2TokenExchange tokenExchange;
+    private final OAuth2Service service;
+    private final ListeningExecutorService responseExecutor;
+
+    @Inject
+    public OAuth2TokenExchangeResource(OAuth2TokenExchange tokenExchange, OAuth2Service service, DispatchExecutor executor)
+    {
+        this.tokenExchange = requireNonNull(tokenExchange, "tokenExchange is null");
+        this.service = requireNonNull(service, "service is null");
+        this.responseExecutor = requireNonNull(executor, "executor is null").getExecutor();
+    }
+
+    @Path("initiate/{authIdHash}")
+    @GET
+    @Produces(MediaType.APPLICATION_JSON)
+    public Response initiateTokenExchange(@PathParam("authIdHash") String authIdHash, @Context HttpServletRequest request)
+    {
+        UriBuilder builder = getSchemeUriBuilder(request);
+        return service.startOAuth2Challenge(builder.build().resolve(CALLBACK_ENDPOINT), Optional.ofNullable(authIdHash));
+    }
+
+    @Path("{authId}")
+    @GET
+    @Produces(MediaType.APPLICATION_JSON)
+    public void getAuthenticationToken(@PathParam("authId") UUID authId, @Suspended AsyncResponse asyncResponse, @Context HttpServletRequest request)
+    {
+        if (authId == null) {
+            throw new BadRequestException();
+        }
+
+        // Do not drop the response from the cache on failure, as this would result in a
+        // hang if the client retries the request. The response will timeout eventually.
+        ListenableFuture tokenFuture = tokenExchange.getTokenPoll(authId);
+        ListenableFuture responseFuture = Futures.transform(tokenFuture, OAuth2TokenExchangeResource::toResponse, responseExecutor);
+        bindAsyncResponse(asyncResponse, responseFuture, responseExecutor)
+                .withTimeout(MAX_POLL_TIME, pendingResponse(request));
+    }
+
+    private static Response toResponse(TokenPoll poll)
+    {
+        if (poll.getError().isPresent()) {
+            return Response.ok(jsonMap("error", poll.getError().get()), APPLICATION_JSON_TYPE).build();
+        }
+        if (poll.getToken().isPresent()) {
+            return Response.ok(jsonMap("token", poll.getToken().get()), APPLICATION_JSON_TYPE).build();
+        }
+        throw new VerifyException("invalid TokenPoll state");
+    }
+
+    private static Response pendingResponse(HttpServletRequest request)
+    {
+        UriBuilder builder = getSchemeUriBuilder(request);
+        return Response.ok(jsonMap("nextUri", builder.build()), APPLICATION_JSON_TYPE).build();
+    }
+
+    @DELETE
+    @Path("{authId}")
+    public Response deleteAuthenticationToken(@PathParam("authId") UUID authId)
+    {
+        if (authId == null) {
+            throw new BadRequestException();
+        }
+
+        tokenExchange.dropToken(authId);
+        return Response
+                .ok()
+                .build();
+    }
+
+    public static String getTokenUri(UUID authId)
+    {
+        return TOKEN_ENDPOINT + authId;
+    }
+
+    public static String getInitiateUri(UUID authId)
+    {
+        return TOKEN_ENDPOINT + "initiate/" + hashAuthId(authId);
+    }
+
+    private static String jsonMap(String key, Object value)
+    {
+        return MAP_CODEC.toJson(ImmutableMap.of(key, value));
+    }
+}
diff --git a/presto-main/src/main/java/com/facebook/presto/server/security/oauth2/OAuth2TokenHandler.java b/presto-main/src/main/java/com/facebook/presto/server/security/oauth2/OAuth2TokenHandler.java
new file mode 100644
index 0000000000000..027d3c27d90f4
--- /dev/null
+++ b/presto-main/src/main/java/com/facebook/presto/server/security/oauth2/OAuth2TokenHandler.java
@@ -0,0 +1,21 @@
+/*
+ * Licensed 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 com.facebook.presto.server.security.oauth2;
+
+public interface OAuth2TokenHandler
+{
+    void setAccessToken(String hashedState, String accessToken);
+
+    void setTokenExchangeError(String hashedState, String errorMessage);
+}
diff --git a/presto-main/src/main/java/com/facebook/presto/server/security/oauth2/OAuth2Utils.java b/presto-main/src/main/java/com/facebook/presto/server/security/oauth2/OAuth2Utils.java
new file mode 100644
index 0000000000000..0b5f8f62db545
--- /dev/null
+++ b/presto-main/src/main/java/com/facebook/presto/server/security/oauth2/OAuth2Utils.java
@@ -0,0 +1,79 @@
+/*
+ * Licensed 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 com.facebook.presto.server.security.oauth2;
+
+import jakarta.servlet.http.HttpServletRequest;
+import jakarta.ws.rs.core.MultivaluedMap;
+import jakarta.ws.rs.core.UriBuilder;
+
+import java.util.List;
+import java.util.Map;
+import java.util.Optional;
+
+import static com.google.common.net.HttpHeaders.X_FORWARDED_PROTO;
+
+public final class OAuth2Utils
+{
+    private OAuth2Utils() {}
+
+    /**
+     * Returns a UriBuilder with the scheme set based upon the X_FORWARDED_PROTO header.
+     * If the header exists on the request we set the scheme to what is in that header. i.e. https.
+     * If the header is not set then we use the scheme on the request.
+     *
+     * Ex: If you are using a load balancer to handle ssl forwarding for Presto. You must set the
+     * X_FORWARDED_PROTO header in the load balancer to 'https'. For any callback or redirect url's
+     * for the OAUTH2 Login flow must use the scheme of https.
+     *
+     * @param request HttpServletRequest
+     * @return a new instance of UriBuilder with the scheme set.
+     */
+    public static UriBuilder getSchemeUriBuilder(HttpServletRequest request)
+    {
+        Optional forwardedProto = Optional.ofNullable(request.getHeader(X_FORWARDED_PROTO));
+
+        UriBuilder builder = UriBuilder.fromUri(getFullRequestURL(request));
+        if (forwardedProto.isPresent()) {
+            builder.scheme(forwardedProto.get());
+        }
+        else {
+            builder.scheme(request.getScheme());
+        }
+
+        return builder;
+    }
+
+    /**
+     * Finds the lastURL query parameter in the request.
+     *
+     * @return Optional the value of the lastURL parameter
+     */
+    public static Optional getLastURLParameter(MultivaluedMap queryParams)
+    {
+        Optional>> lastUrl = queryParams.entrySet().stream().filter(qp -> qp.getKey().equals("lastURL")).findFirst();
+        if (lastUrl.isPresent() && lastUrl.get().getValue().size() > 0) {
+            return Optional.ofNullable(lastUrl.get().getValue().get(0));
+        }
+
+        return Optional.empty();
+    }
+
+    public static String getFullRequestURL(HttpServletRequest request)
+    {
+        StringBuilder requestURL = new StringBuilder(request.getRequestURL());
+        String queryString = request.getQueryString();
+
+        return queryString == null ? requestURL.toString() : requestURL.append("?").append(queryString).toString();
+    }
+}
diff --git a/presto-main/src/main/java/com/facebook/presto/server/security/oauth2/OAuth2WebUiAuthenticationManager.java b/presto-main/src/main/java/com/facebook/presto/server/security/oauth2/OAuth2WebUiAuthenticationManager.java
new file mode 100644
index 0000000000000..925019f5ac442
--- /dev/null
+++ b/presto-main/src/main/java/com/facebook/presto/server/security/oauth2/OAuth2WebUiAuthenticationManager.java
@@ -0,0 +1,120 @@
+/*
+ * Licensed 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 com.facebook.presto.server.security.oauth2;
+
+import com.facebook.airlift.http.server.AuthenticationException;
+import com.facebook.airlift.log.Logger;
+import com.facebook.presto.server.security.WebUiAuthenticationManager;
+import com.facebook.presto.server.security.oauth2.TokenPairSerializer.TokenPair;
+import jakarta.servlet.FilterChain;
+import jakarta.servlet.ServletException;
+import jakarta.servlet.http.Cookie;
+import jakarta.servlet.http.HttpServletRequest;
+import jakarta.servlet.http.HttpServletResponse;
+import jakarta.ws.rs.core.UriBuilder;
+
+import javax.inject.Inject;
+
+import java.io.IOException;
+import java.security.Principal;
+import java.time.Duration;
+import java.time.Instant;
+import java.util.Optional;
+
+import static com.facebook.presto.server.security.AuthenticationFilter.withPrincipal;
+import static com.facebook.presto.server.security.oauth2.OAuth2Authenticator.extractTokenFromCookie;
+import static com.facebook.presto.server.security.oauth2.OAuth2CallbackResource.CALLBACK_ENDPOINT;
+import static com.facebook.presto.server.security.oauth2.OAuth2Utils.getSchemeUriBuilder;
+import static java.util.Objects.requireNonNull;
+
+public class OAuth2WebUiAuthenticationManager
+        implements WebUiAuthenticationManager
+{
+    private static final Logger logger = Logger.get(OAuth2WebUiAuthenticationManager.class);
+    private final OAuth2Service oAuth2Service;
+    private final OAuth2Authenticator oAuth2Authenticator;
+    private final TokenPairSerializer tokenPairSerializer;
+    private final OAuth2Client client;
+    private final Optional tokenExpiration;
+
+    @Inject
+    public OAuth2WebUiAuthenticationManager(OAuth2Service oAuth2Service, OAuth2Authenticator oAuth2Authenticator, TokenPairSerializer tokenPairSerializer, OAuth2Client client, @ForRefreshTokens Optional tokenExpiration)
+    {
+        this.oAuth2Service = requireNonNull(oAuth2Service, "oauth2Service is null");
+        this.oAuth2Authenticator = requireNonNull(oAuth2Authenticator, "oauth2Authenticator is null");
+        this.tokenPairSerializer = requireNonNull(tokenPairSerializer, "tokenPairSerializer is null");
+        this.client = requireNonNull(client, "oauth2Client is null");
+        this.tokenExpiration = requireNonNull(tokenExpiration, "tokenExpiration is null");
+    }
+
+    public void handleRequest(HttpServletRequest request, HttpServletResponse response, FilterChain nextFilter)
+            throws IOException, ServletException
+    {
+        try {
+            Principal principal = this.oAuth2Authenticator.authenticate(request);
+            nextFilter.doFilter(withPrincipal(request, principal), response);
+        }
+        catch (AuthenticationException e) {
+            needAuthentication(request, response);
+        }
+    }
+
+    private Optional getTokenPair(HttpServletRequest request)
+    {
+        try {
+            Optional token = extractTokenFromCookie(request);
+            if (token.isPresent()) {
+                return Optional.ofNullable(tokenPairSerializer.deserialize(token.get()));
+            }
+            else {
+                return Optional.empty();
+            }
+        }
+        catch (Exception e) {
+            logger.error(e, "Exception occurred during token pair deserialization");
+            return Optional.empty();
+        }
+    }
+
+    private void needAuthentication(HttpServletRequest request, HttpServletResponse response)
+            throws IOException
+    {
+        Optional tokenPair = getTokenPair(request);
+        Optional refreshToken = tokenPair.flatMap(TokenPair::getRefreshToken);
+        if (refreshToken.isPresent()) {
+            try {
+                OAuth2Client.Response refreshRes = client.refreshTokens(refreshToken.get());
+                String serializeToken = tokenPairSerializer.serialize(TokenPair.fromOAuth2Response(refreshRes));
+                UriBuilder builder = getSchemeUriBuilder(request);
+                Cookie newCookie = NonceCookie.toServletCookie(OAuthWebUiCookie.create(serializeToken, tokenExpiration.map(expiration -> Instant.now().plus(expiration)).orElse(refreshRes.getExpiration())));
+                response.addCookie(newCookie);
+                response.sendRedirect(builder.build().toString());
+            }
+            catch (ChallengeFailedException e) {
+                logger.error(e, "Token refresh challenge has failed");
+                this.startOauth2Challenge(request, response);
+            }
+        }
+        else {
+            this.startOauth2Challenge(request, response);
+        }
+    }
+
+    private void startOauth2Challenge(HttpServletRequest request, HttpServletResponse response)
+            throws IOException
+    {
+        UriBuilder builder = getSchemeUriBuilder(request);
+        this.oAuth2Service.startOAuth2Challenge(builder.build().resolve(CALLBACK_ENDPOINT), Optional.empty(), response);
+    }
+}
diff --git a/presto-main/src/main/java/com/facebook/presto/server/security/oauth2/OAuthWebUiCookie.java b/presto-main/src/main/java/com/facebook/presto/server/security/oauth2/OAuthWebUiCookie.java
new file mode 100644
index 0000000000000..867e7b60c3dd8
--- /dev/null
+++ b/presto-main/src/main/java/com/facebook/presto/server/security/oauth2/OAuthWebUiCookie.java
@@ -0,0 +1,61 @@
+/*
+ * Licensed 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 com.facebook.presto.server.security.oauth2;
+
+import jakarta.ws.rs.core.NewCookie;
+
+import java.time.Instant;
+import java.util.Date;
+
+import static jakarta.ws.rs.core.Cookie.DEFAULT_VERSION;
+import static jakarta.ws.rs.core.NewCookie.DEFAULT_MAX_AGE;
+
+public final class OAuthWebUiCookie
+{
+    // prefix according to: https://tools.ietf.org/html/draft-ietf-httpbis-rfc6265bis-05#section-4.1.3.1
+    public static final String OAUTH2_COOKIE = "__Secure-Presto-OAuth2-Token";
+
+    public static final String API_PATH = "/";
+
+    private OAuthWebUiCookie() {}
+
+    public static NewCookie create(String token, Instant tokenExpiration)
+    {
+        return new NewCookie(
+                OAUTH2_COOKIE,
+                token,
+                API_PATH,
+                null,
+                DEFAULT_VERSION,
+                null,
+                DEFAULT_MAX_AGE,
+                Date.from(tokenExpiration),
+                true,
+                true);
+    }
+    public static NewCookie delete()
+    {
+        return new NewCookie(
+                OAUTH2_COOKIE,
+                "delete",
+                API_PATH,
+                null,
+                DEFAULT_VERSION,
+                null,
+                0,
+                null,
+                true,
+                true);
+    }
+}
diff --git a/presto-main/src/main/java/com/facebook/presto/server/security/oauth2/OidcDiscovery.java b/presto-main/src/main/java/com/facebook/presto/server/security/oauth2/OidcDiscovery.java
new file mode 100644
index 0000000000000..1b6bae30cc8c5
--- /dev/null
+++ b/presto-main/src/main/java/com/facebook/presto/server/security/oauth2/OidcDiscovery.java
@@ -0,0 +1,159 @@
+/*
+ * Licensed 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 com.facebook.presto.server.security.oauth2;
+
+import com.facebook.airlift.json.JsonObjectMapperProvider;
+import com.facebook.airlift.log.Logger;
+import com.fasterxml.jackson.core.JsonProcessingException;
+import com.fasterxml.jackson.databind.JsonNode;
+import com.fasterxml.jackson.databind.ObjectMapper;
+import com.nimbusds.oauth2.sdk.ParseException;
+import com.nimbusds.oauth2.sdk.http.HTTPResponse;
+import com.nimbusds.oauth2.sdk.id.Issuer;
+import com.nimbusds.openid.connect.sdk.op.OIDCProviderConfigurationRequest;
+import com.nimbusds.openid.connect.sdk.op.OIDCProviderMetadata;
+import net.jodah.failsafe.Failsafe;
+import net.jodah.failsafe.RetryPolicy;
+
+import javax.inject.Inject;
+
+import java.net.URI;
+import java.time.Duration;
+import java.util.Optional;
+
+import static com.facebook.airlift.http.client.HttpStatus.OK;
+import static com.facebook.airlift.http.client.HttpStatus.REQUEST_TIMEOUT;
+import static com.facebook.airlift.http.client.HttpStatus.TOO_MANY_REQUESTS;
+import static com.facebook.presto.server.security.oauth2.StaticOAuth2ServerConfiguration.ACCESS_TOKEN_ISSUER;
+import static com.facebook.presto.server.security.oauth2.StaticOAuth2ServerConfiguration.AUTH_URL;
+import static com.facebook.presto.server.security.oauth2.StaticOAuth2ServerConfiguration.JWKS_URL;
+import static com.facebook.presto.server.security.oauth2.StaticOAuth2ServerConfiguration.TOKEN_URL;
+import static com.facebook.presto.server.security.oauth2.StaticOAuth2ServerConfiguration.USERINFO_URL;
+import static com.google.common.base.Preconditions.checkState;
+import static java.lang.String.format;
+import static java.util.Objects.requireNonNull;
+
+public class OidcDiscovery
+        implements OAuth2ServerConfigProvider
+{
+    private static final Logger LOG = Logger.get(OidcDiscovery.class);
+
+    private static final ObjectMapper OBJECT_MAPPER = new JsonObjectMapperProvider().get();
+    private final Issuer issuer;
+    private final Duration discoveryTimeout;
+    private final boolean userinfoEndpointEnabled;
+    private final Optional accessTokenIssuer;
+    private final Optional authUrl;
+    private final Optional tokenUrl;
+    private final Optional jwksUrl;
+    private final Optional userinfoUrl;
+    private final NimbusHttpClient httpClient;
+
+    @Inject
+    public OidcDiscovery(OAuth2Config oauthConfig, OidcDiscoveryConfig oidcConfig, NimbusHttpClient httpClient)
+    {
+        requireNonNull(oauthConfig, "oauthConfig is null");
+        issuer = new Issuer(requireNonNull(oauthConfig.getIssuer(), "issuer is null"));
+        requireNonNull(oidcConfig, "oidcConfig is null");
+        userinfoEndpointEnabled = oidcConfig.isUserinfoEndpointEnabled();
+        discoveryTimeout = Duration.ofMillis(requireNonNull(oidcConfig.getDiscoveryTimeout(), "discoveryTimeout is null").toMillis());
+        accessTokenIssuer = requireNonNull(oidcConfig.getAccessTokenIssuer(), "accessTokenIssuer is null");
+        authUrl = requireNonNull(oidcConfig.getAuthUrl(), "authUrl is null");
+        tokenUrl = requireNonNull(oidcConfig.getTokenUrl(), "tokenUrl is null");
+        jwksUrl = requireNonNull(oidcConfig.getJwksUrl(), "jwksUrl is null");
+        userinfoUrl = requireNonNull(oidcConfig.getUserinfoUrl(), "userinfoUrl is null");
+        this.httpClient = requireNonNull(httpClient, "httpClient is null");
+    }
+
+    @Override
+    public OAuth2ServerConfig get()
+    {
+        return Failsafe.with(new RetryPolicy<>()
+                        .withMaxAttempts(-1)
+                        .withMaxDuration(discoveryTimeout)
+                        .withDelay(Duration.ofSeconds(1))
+                        .abortOn(IllegalStateException.class)
+                        .onFailedAttempt(attempt -> LOG.debug("OpenID Connect Metadata read failed: %s", attempt.getLastFailure())))
+                .get(() -> httpClient.execute(new OIDCProviderConfigurationRequest(issuer), this::parseConfigurationResponse));
+    }
+
+    private OAuth2ServerConfig parseConfigurationResponse(HTTPResponse response)
+            throws ParseException
+    {
+        int statusCode = response.getStatusCode();
+        if (statusCode != OK.code()) {
+            // stop on any client errors other than REQUEST_TIMEOUT and TOO_MANY_REQUESTS
+            if (statusCode < 400 || statusCode >= 500 || statusCode == REQUEST_TIMEOUT.code() || statusCode == TOO_MANY_REQUESTS.code()) {
+                throw new RuntimeException("Invalid response from OpenID Metadata endpoint: " + statusCode);
+            }
+            else {
+                throw new IllegalStateException(format("Invalid response from OpenID Metadata endpoint. Expected response code to be %s, but was %s", OK.code(), statusCode));
+            }
+        }
+        return readConfiguration(response.getContent());
+    }
+
+    private OAuth2ServerConfig readConfiguration(String body)
+            throws ParseException
+    {
+        OIDCProviderMetadata metadata = OIDCProviderMetadata.parse(body);
+        checkMetadataState(issuer.equals(metadata.getIssuer()), "The value of the \"issuer\" claim in Metadata document different than the Issuer URL used for the Configuration Request.");
+        try {
+            JsonNode metadataJson = OBJECT_MAPPER.readTree(body);
+            Optional userinfoEndpoint;
+            if (userinfoEndpointEnabled) {
+                userinfoEndpoint = getOptionalField("userinfo_endpoint", Optional.ofNullable(metadata.getUserInfoEndpointURI()).map(URI::toString), USERINFO_URL, userinfoUrl);
+            }
+            else {
+                userinfoEndpoint = Optional.empty();
+            }
+            return new OAuth2ServerConfig(
+                    // AD FS server can include "access_token_issuer" field in OpenID Provider Metadata.
+                    // It's not a part of the OIDC standard thus have to be handled separately.
+                    // see: https://docs.microsoft.com/en-us/openspecs/windows_protocols/ms-oidce/f629647a-4825-465b-80bb-32c7e9cec2c8
+                    getOptionalField("access_token_issuer", Optional.ofNullable(metadataJson.get("access_token_issuer")).map(JsonNode::textValue), ACCESS_TOKEN_ISSUER, accessTokenIssuer),
+                    getRequiredField("authorization_endpoint", metadata.getAuthorizationEndpointURI(), AUTH_URL, authUrl),
+                    getRequiredField("token_endpoint", metadata.getTokenEndpointURI(), TOKEN_URL, tokenUrl),
+                    getRequiredField("jwks_uri", metadata.getJWKSetURI(), JWKS_URL, jwksUrl),
+                    userinfoEndpoint.map(URI::create));
+        }
+        catch (JsonProcessingException e) {
+            throw new ParseException("Invalid JSON value", e);
+        }
+    }
+
+    private static URI getRequiredField(String metadataField, URI metadataValue, String configurationField, Optional configurationValue)
+    {
+        Optional uri = getOptionalField(metadataField, Optional.ofNullable(metadataValue).map(URI::toString), configurationField, configurationValue);
+        checkMetadataState(uri.isPresent(), "Missing required \"%s\" property.", metadataField);
+        return URI.create(uri.get());
+    }
+
+    private static Optional getOptionalField(String metadataField, Optional metadataValue, String configurationField, Optional configurationValue)
+    {
+        if (configurationValue.isPresent()) {
+            if (!configurationValue.equals(metadataValue)) {
+                LOG.warn("Overriding \"%s=%s\" from OpenID metadata document with value \"%s=%s\" defined in configuration",
+                        metadataField, metadataValue.orElse(""), configurationField, configurationValue.orElse(""));
+            }
+            return configurationValue;
+        }
+        return metadataValue;
+    }
+
+    private static void checkMetadataState(boolean expression, String additionalMessage, String... additionalMessageArgs)
+    {
+        checkState(expression, "Invalid response from OpenID Metadata endpoint. " + additionalMessage, (Object[]) additionalMessageArgs);
+    }
+}
diff --git a/presto-main/src/main/java/com/facebook/presto/server/security/oauth2/OidcDiscoveryConfig.java b/presto-main/src/main/java/com/facebook/presto/server/security/oauth2/OidcDiscoveryConfig.java
new file mode 100644
index 0000000000000..d5d29589fc5a4
--- /dev/null
+++ b/presto-main/src/main/java/com/facebook/presto/server/security/oauth2/OidcDiscoveryConfig.java
@@ -0,0 +1,148 @@
+/*
+ * Licensed 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 com.facebook.presto.server.security.oauth2;
+
+import com.facebook.airlift.configuration.Config;
+import com.facebook.airlift.configuration.ConfigDescription;
+import com.facebook.airlift.units.Duration;
+import jakarta.validation.constraints.NotNull;
+
+import java.util.Optional;
+
+import static com.facebook.presto.server.security.oauth2.StaticOAuth2ServerConfiguration.ACCESS_TOKEN_ISSUER;
+import static com.facebook.presto.server.security.oauth2.StaticOAuth2ServerConfiguration.AUTH_URL;
+import static com.facebook.presto.server.security.oauth2.StaticOAuth2ServerConfiguration.JWKS_URL;
+import static com.facebook.presto.server.security.oauth2.StaticOAuth2ServerConfiguration.TOKEN_URL;
+import static com.facebook.presto.server.security.oauth2.StaticOAuth2ServerConfiguration.USERINFO_URL;
+import static java.util.concurrent.TimeUnit.SECONDS;
+
+public class OidcDiscoveryConfig
+{
+    private Duration discoveryTimeout = new Duration(30, SECONDS);
+    private boolean userinfoEndpointEnabled = true;
+
+    //TODO Left for backward compatibility, remove after the next release/a couple of releases
+    private Optional accessTokenIssuer = Optional.empty();
+    private Optional authUrl = Optional.empty();
+    private Optional tokenUrl = Optional.empty();
+    private Optional jwksUrl = Optional.empty();
+    private Optional userinfoUrl = Optional.empty();
+
+    @NotNull
+    public Duration getDiscoveryTimeout()
+    {
+        return discoveryTimeout;
+    }
+
+    @Config("http-server.authentication.oauth2.oidc.discovery.timeout")
+    @ConfigDescription("OpenID Connect discovery timeout")
+    public OidcDiscoveryConfig setDiscoveryTimeout(Duration discoveryTimeout)
+    {
+        this.discoveryTimeout = discoveryTimeout;
+        return this;
+    }
+
+    public boolean isUserinfoEndpointEnabled()
+    {
+        return userinfoEndpointEnabled;
+    }
+
+    @Config("http-server.authentication.oauth2.oidc.use-userinfo-endpoint")
+    @ConfigDescription("Use userinfo endpoint from OpenID connect metadata document")
+    public OidcDiscoveryConfig setUserinfoEndpointEnabled(boolean userinfoEndpointEnabled)
+    {
+        this.userinfoEndpointEnabled = userinfoEndpointEnabled;
+        return this;
+    }
+
+    @NotNull
+    @Deprecated
+    public Optional getAccessTokenIssuer()
+    {
+        return accessTokenIssuer;
+    }
+
+    @Config(ACCESS_TOKEN_ISSUER)
+    @ConfigDescription("The required issuer for access tokens")
+    @Deprecated
+    public OidcDiscoveryConfig setAccessTokenIssuer(String accessTokenIssuer)
+    {
+        this.accessTokenIssuer = Optional.ofNullable(accessTokenIssuer);
+        return this;
+    }
+
+    @NotNull
+    @Deprecated
+    public Optional getAuthUrl()
+    {
+        return authUrl;
+    }
+
+    @Config(AUTH_URL)
+    @ConfigDescription("URL of the authorization server's authorization endpoint")
+    @Deprecated
+    public OidcDiscoveryConfig setAuthUrl(String authUrl)
+    {
+        this.authUrl = Optional.ofNullable(authUrl);
+        return this;
+    }
+
+    @NotNull
+    @Deprecated
+    public Optional getTokenUrl()
+    {
+        return tokenUrl;
+    }
+
+    @Config(TOKEN_URL)
+    @ConfigDescription("URL of the authorization server's token endpoint")
+    @Deprecated
+    public OidcDiscoveryConfig setTokenUrl(String tokenUrl)
+    {
+        this.tokenUrl = Optional.ofNullable(tokenUrl);
+        return this;
+    }
+
+    @NotNull
+    @Deprecated
+    public Optional getJwksUrl()
+    {
+        return jwksUrl;
+    }
+
+    @Config(JWKS_URL)
+    @ConfigDescription("URL of the authorization server's JWKS (JSON Web Key Set) endpoint")
+    @Deprecated
+    public OidcDiscoveryConfig setJwksUrl(String jwksUrl)
+    {
+        this.jwksUrl = Optional.ofNullable(jwksUrl);
+        return this;
+    }
+
+    @NotNull
+    @Deprecated
+    public Optional getUserinfoUrl()
+    {
+        return userinfoUrl;
+    }
+
+    @Config(USERINFO_URL)
+    @ConfigDescription("URL of the userinfo endpoint")
+    @Deprecated
+    public OidcDiscoveryConfig setUserinfoUrl(String userinfoUrl)
+    {
+        this.userinfoUrl = Optional.ofNullable(userinfoUrl);
+        return this;
+    }
+}
diff --git a/presto-main/src/main/java/com/facebook/presto/server/security/oauth2/RefreshTokensConfig.java b/presto-main/src/main/java/com/facebook/presto/server/security/oauth2/RefreshTokensConfig.java
new file mode 100644
index 0000000000000..7d25c142f0332
--- /dev/null
+++ b/presto-main/src/main/java/com/facebook/presto/server/security/oauth2/RefreshTokensConfig.java
@@ -0,0 +1,96 @@
+/*
+ * Licensed 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 com.facebook.presto.server.security.oauth2;
+
+import com.facebook.airlift.configuration.Config;
+import com.facebook.airlift.configuration.ConfigDescription;
+import com.facebook.airlift.configuration.ConfigSecuritySensitive;
+import com.facebook.airlift.units.Duration;
+import io.jsonwebtoken.io.Decoders;
+import io.jsonwebtoken.security.Keys;
+import jakarta.validation.constraints.NotEmpty;
+
+import javax.crypto.SecretKey;
+
+import static com.google.common.base.Strings.isNullOrEmpty;
+import static java.util.concurrent.TimeUnit.HOURS;
+
+public class RefreshTokensConfig
+{
+    private Duration tokenExpiration = Duration.succinctDuration(1, HOURS);
+    private static final String coordinator = "Presto_coordinator";
+    private String issuer = coordinator;
+    private String audience = coordinator;
+
+    private SecretKey secretKey;
+
+    public Duration getTokenExpiration()
+    {
+        return tokenExpiration;
+    }
+
+    @Config("http-server.authentication.oauth2.refresh-tokens.issued-token.timeout")
+    @ConfigDescription("Expiration time for issued token. It needs to be equal or lower than duration of refresh token issued by IdP")
+    public RefreshTokensConfig setTokenExpiration(Duration tokenExpiration)
+    {
+        this.tokenExpiration = tokenExpiration;
+        return this;
+    }
+
+    @NotEmpty
+    public String getIssuer()
+    {
+        return issuer;
+    }
+
+    @Config("http-server.authentication.oauth2.refresh-tokens.issued-token.issuer")
+    @ConfigDescription("Issuer representing this coordinator instance, that will be used in issued token. In addition current Version will be added to it")
+    public RefreshTokensConfig setIssuer(String issuer)
+    {
+        this.issuer = issuer;
+        return this;
+    }
+
+    @NotEmpty
+    public String getAudience()
+    {
+        return audience;
+    }
+
+    @Config("http-server.authentication.oauth2.refresh-tokens.issued-token.audience")
+    @ConfigDescription("Audience representing this coordinator instance, that will be used in issued token")
+    public RefreshTokensConfig setAudience(String audience)
+    {
+        this.audience = audience;
+        return this;
+    }
+
+    @Config("http-server.authentication.oauth2.refresh-tokens.secret-key")
+    @ConfigDescription("Base64 encoded secret key used to encrypt generated token")
+    @ConfigSecuritySensitive
+    public RefreshTokensConfig setSecretKey(String key)
+    {
+        if (isNullOrEmpty(key)) {
+            return this;
+        }
+
+        secretKey = Keys.hmacShaKeyFor(Decoders.BASE64.decode(key));
+        return this;
+    }
+
+    public SecretKey getSecretKey()
+    {
+        return secretKey;
+    }
+}
diff --git a/presto-main/src/main/java/com/facebook/presto/server/security/oauth2/StaticConfigurationProvider.java b/presto-main/src/main/java/com/facebook/presto/server/security/oauth2/StaticConfigurationProvider.java
new file mode 100644
index 0000000000000..627c42b3a5814
--- /dev/null
+++ b/presto-main/src/main/java/com/facebook/presto/server/security/oauth2/StaticConfigurationProvider.java
@@ -0,0 +1,44 @@
+/*
+ * Licensed 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 com.facebook.presto.server.security.oauth2;
+
+import javax.inject.Inject;
+
+import java.net.URI;
+
+import static java.util.Objects.requireNonNull;
+
+public class StaticConfigurationProvider
+        implements OAuth2ServerConfigProvider
+{
+    private final OAuth2ServerConfig config;
+
+    @Inject
+    StaticConfigurationProvider(StaticOAuth2ServerConfiguration config)
+    {
+        requireNonNull(config, "config is null");
+        this.config = new OAuth2ServerConfig(
+                config.getAccessTokenIssuer(),
+                URI.create(config.getAuthUrl()),
+                URI.create(config.getTokenUrl()),
+                URI.create(config.getJwksUrl()),
+                config.getUserinfoUrl().map(URI::create));
+    }
+
+    @Override
+    public OAuth2ServerConfig get()
+    {
+        return config;
+    }
+}
diff --git a/presto-main/src/main/java/com/facebook/presto/server/security/oauth2/StaticOAuth2ServerConfiguration.java b/presto-main/src/main/java/com/facebook/presto/server/security/oauth2/StaticOAuth2ServerConfiguration.java
new file mode 100644
index 0000000000000..51f116abe09b5
--- /dev/null
+++ b/presto-main/src/main/java/com/facebook/presto/server/security/oauth2/StaticOAuth2ServerConfiguration.java
@@ -0,0 +1,104 @@
+/*
+ * Licensed 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 com.facebook.presto.server.security.oauth2;
+
+import com.facebook.airlift.configuration.Config;
+import com.facebook.airlift.configuration.ConfigDescription;
+import jakarta.validation.constraints.NotNull;
+
+import java.util.Optional;
+
+public class StaticOAuth2ServerConfiguration
+{
+    public static final String ACCESS_TOKEN_ISSUER = "http-server.authentication.oauth2.access-token-issuer";
+    public static final String AUTH_URL = "http-server.authentication.oauth2.auth-url";
+    public static final String TOKEN_URL = "http-server.authentication.oauth2.token-url";
+    public static final String JWKS_URL = "http-server.authentication.oauth2.jwks-url";
+    public static final String USERINFO_URL = "http-server.authentication.oauth2.userinfo-url";
+
+    private Optional accessTokenIssuer = Optional.empty();
+    private String authUrl;
+    private String tokenUrl;
+    private String jwksUrl;
+    private Optional userinfoUrl = Optional.empty();
+
+    @NotNull
+    public Optional getAccessTokenIssuer()
+    {
+        return accessTokenIssuer;
+    }
+
+    @Config(ACCESS_TOKEN_ISSUER)
+    @ConfigDescription("The required issuer for access tokens")
+    public StaticOAuth2ServerConfiguration setAccessTokenIssuer(String accessTokenIssuer)
+    {
+        this.accessTokenIssuer = Optional.ofNullable(accessTokenIssuer);
+        return this;
+    }
+
+    @NotNull
+    public String getAuthUrl()
+    {
+        return authUrl;
+    }
+
+    @Config(AUTH_URL)
+    @ConfigDescription("URL of the authorization server's authorization endpoint")
+    public StaticOAuth2ServerConfiguration setAuthUrl(String authUrl)
+    {
+        this.authUrl = authUrl;
+        return this;
+    }
+
+    @NotNull
+    public String getTokenUrl()
+    {
+        return tokenUrl;
+    }
+
+    @Config(TOKEN_URL)
+    @ConfigDescription("URL of the authorization server's token endpoint")
+    public StaticOAuth2ServerConfiguration setTokenUrl(String tokenUrl)
+    {
+        this.tokenUrl = tokenUrl;
+        return this;
+    }
+
+    @NotNull
+    public String getJwksUrl()
+    {
+        return jwksUrl;
+    }
+
+    @Config(JWKS_URL)
+    @ConfigDescription("URL of the authorization server's JWKS (JSON Web Key Set) endpoint")
+    public StaticOAuth2ServerConfiguration setJwksUrl(String jwksUrl)
+    {
+        this.jwksUrl = jwksUrl;
+        return this;
+    }
+
+    public Optional getUserinfoUrl()
+    {
+        return userinfoUrl;
+    }
+
+    @Config(USERINFO_URL)
+    @ConfigDescription("URL of the userinfo endpoint")
+    public StaticOAuth2ServerConfiguration setUserinfoUrl(String userinfoUrl)
+    {
+        this.userinfoUrl = Optional.ofNullable(userinfoUrl);
+        return this;
+    }
+}
diff --git a/presto-main/src/main/java/com/facebook/presto/server/security/oauth2/TokenPairSerializer.java b/presto-main/src/main/java/com/facebook/presto/server/security/oauth2/TokenPairSerializer.java
new file mode 100644
index 0000000000000..f178a4048adc8
--- /dev/null
+++ b/presto-main/src/main/java/com/facebook/presto/server/security/oauth2/TokenPairSerializer.java
@@ -0,0 +1,91 @@
+/*
+ * Licensed 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 com.facebook.presto.server.security.oauth2;
+
+import com.facebook.presto.server.security.oauth2.OAuth2Client.Response;
+import jakarta.annotation.Nullable;
+
+import java.util.Date;
+import java.util.Optional;
+
+import static java.lang.Long.MAX_VALUE;
+import static java.util.Objects.requireNonNull;
+
+public interface TokenPairSerializer
+{
+    TokenPairSerializer ACCESS_TOKEN_ONLY_SERIALIZER = new TokenPairSerializer()
+    {
+        @Override
+        public TokenPair deserialize(String token)
+        {
+            return TokenPair.accessToken(token);
+        }
+
+        @Override
+        public String serialize(TokenPair tokenPair)
+        {
+            return tokenPair.getAccessToken();
+        }
+    };
+
+    TokenPair deserialize(String token);
+
+    String serialize(TokenPair tokenPair);
+
+    class TokenPair
+    {
+        private final String accessToken;
+        private final Date expiration;
+        private final Optional refreshToken;
+
+        private TokenPair(String accessToken, Date expiration, Optional refreshToken)
+        {
+            this.accessToken = requireNonNull(accessToken, "accessToken is nul");
+            this.expiration = requireNonNull(expiration, "expiration is null");
+            this.refreshToken = requireNonNull(refreshToken, "refreshToken is null");
+        }
+
+        public static TokenPair accessToken(String accessToken)
+        {
+            return new TokenPair(accessToken, new Date(MAX_VALUE), Optional.empty());
+        }
+
+        public static TokenPair fromOAuth2Response(Response tokens)
+        {
+            requireNonNull(tokens, "tokens is null");
+            return new TokenPair(tokens.getAccessToken(), Date.from(tokens.getExpiration()), tokens.getRefreshToken());
+        }
+
+        public static TokenPair accessAndRefreshTokens(String accessToken, Date expiration, @Nullable String refreshToken)
+        {
+            return new TokenPair(accessToken, expiration, Optional.ofNullable(refreshToken));
+        }
+
+        public String getAccessToken()
+        {
+            return accessToken;
+        }
+
+        public Date getExpiration()
+        {
+            return expiration;
+        }
+
+        public Optional getRefreshToken()
+        {
+            return refreshToken;
+        }
+    }
+}
diff --git a/presto-main/src/main/java/com/facebook/presto/server/security/oauth2/TokenRefresher.java b/presto-main/src/main/java/com/facebook/presto/server/security/oauth2/TokenRefresher.java
new file mode 100644
index 0000000000000..411103f6f8724
--- /dev/null
+++ b/presto-main/src/main/java/com/facebook/presto/server/security/oauth2/TokenRefresher.java
@@ -0,0 +1,68 @@
+/*
+ * Licensed 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 com.facebook.presto.server.security.oauth2;
+
+import com.facebook.presto.server.security.oauth2.OAuth2Client.Response;
+import com.facebook.presto.server.security.oauth2.TokenPairSerializer.TokenPair;
+
+import java.util.Optional;
+import java.util.UUID;
+
+import static com.facebook.presto.server.security.oauth2.OAuth2TokenExchange.hashAuthId;
+import static java.util.Objects.requireNonNull;
+
+public class TokenRefresher
+{
+    private final TokenPairSerializer tokenAssembler;
+    private final OAuth2TokenHandler tokenHandler;
+    private final OAuth2Client client;
+
+    public TokenRefresher(TokenPairSerializer tokenAssembler, OAuth2TokenHandler tokenHandler, OAuth2Client client)
+    {
+        this.tokenAssembler = requireNonNull(tokenAssembler, "tokenAssembler is null");
+        this.tokenHandler = requireNonNull(tokenHandler, "tokenHandler is null");
+        this.client = requireNonNull(client, "oAuth2Client is null");
+    }
+
+    public Optional refreshToken(TokenPair tokenPair)
+    {
+        requireNonNull(tokenPair, "tokenPair is null");
+
+        Optional refreshToken = tokenPair.getRefreshToken();
+        if (refreshToken.isPresent()) {
+            UUID refreshingId = UUID.randomUUID();
+            try {
+                refreshToken(refreshToken.get(), refreshingId);
+                return Optional.of(refreshingId);
+            }
+            // If Refresh token has expired then restart the flow
+            catch (RuntimeException exception) {
+                return Optional.empty();
+            }
+        }
+        return Optional.empty();
+    }
+
+    private void refreshToken(String refreshToken, UUID refreshingId)
+    {
+        try {
+            Response response = client.refreshTokens(refreshToken);
+            String serializedToken = tokenAssembler.serialize(TokenPair.fromOAuth2Response(response));
+            tokenHandler.setAccessToken(hashAuthId(refreshingId), serializedToken);
+        }
+        catch (ChallengeFailedException e) {
+            tokenHandler.setTokenExchangeError(hashAuthId(refreshingId), "Token refreshing has failed: " + e.getMessage());
+        }
+    }
+}
diff --git a/presto-main/src/main/java/com/facebook/presto/server/security/oauth2/ZstdCodec.java b/presto-main/src/main/java/com/facebook/presto/server/security/oauth2/ZstdCodec.java
new file mode 100644
index 0000000000000..1065d90416240
--- /dev/null
+++ b/presto-main/src/main/java/com/facebook/presto/server/security/oauth2/ZstdCodec.java
@@ -0,0 +1,53 @@
+/*
+ * Licensed 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 com.facebook.presto.server.security.oauth2;
+
+import io.airlift.compress.zstd.ZstdCompressor;
+import io.airlift.compress.zstd.ZstdDecompressor;
+import io.jsonwebtoken.CompressionCodec;
+import io.jsonwebtoken.CompressionException;
+
+import static java.lang.Math.toIntExact;
+import static java.util.Arrays.copyOfRange;
+
+public class ZstdCodec
+        implements CompressionCodec
+{
+    public static final String CODEC_NAME = "ZSTD";
+
+    @Override
+    public String getAlgorithmName()
+    {
+        return CODEC_NAME;
+    }
+
+    @Override
+    public byte[] compress(byte[] bytes)
+            throws CompressionException
+    {
+        ZstdCompressor compressor = new ZstdCompressor();
+        byte[] compressed = new byte[compressor.maxCompressedLength(bytes.length)];
+        int outputSize = compressor.compress(bytes, 0, bytes.length, compressed, 0, compressed.length);
+        return copyOfRange(compressed, 0, outputSize);
+    }
+
+    @Override
+    public byte[] decompress(byte[] bytes)
+            throws CompressionException
+    {
+        byte[] output = new byte[toIntExact(ZstdDecompressor.getDecompressedSize(bytes, 0, bytes.length))];
+        new ZstdDecompressor().decompress(bytes, 0, bytes.length, output, 0, output.length);
+        return output;
+    }
+}
diff --git a/presto-main/src/main/java/com/facebook/presto/server/testing/TestingPrestoServer.java b/presto-main/src/main/java/com/facebook/presto/server/testing/TestingPrestoServer.java
index 1cb04591b15da..980ba54f389b8 100644
--- a/presto-main/src/main/java/com/facebook/presto/server/testing/TestingPrestoServer.java
+++ b/presto-main/src/main/java/com/facebook/presto/server/testing/TestingPrestoServer.java
@@ -34,6 +34,7 @@
 import com.facebook.drift.transport.netty.server.DriftNettyServerTransport;
 import com.facebook.presto.ClientRequestFilterManager;
 import com.facebook.presto.ClientRequestFilterModule;
+import com.facebook.presto.builtin.tools.WorkerFunctionRegistryTool;
 import com.facebook.presto.connector.ConnectorManager;
 import com.facebook.presto.cost.StatsCalculator;
 import com.facebook.presto.dispatcher.DispatchManager;
@@ -51,6 +52,7 @@
 import com.facebook.presto.memory.LocalMemoryManager;
 import com.facebook.presto.metadata.AllNodes;
 import com.facebook.presto.metadata.CatalogManager;
+import com.facebook.presto.metadata.FunctionAndTypeManager;
 import com.facebook.presto.metadata.InternalNode;
 import com.facebook.presto.metadata.InternalNodeManager;
 import com.facebook.presto.metadata.Metadata;
@@ -147,6 +149,8 @@ public class TestingPrestoServer
     private final boolean preserveData;
     private final LifeCycleManager lifeCycleManager;
     private final PluginManager pluginManager;
+    private final FunctionAndTypeManager functionAndTypeManager;
+    private final WorkerFunctionRegistryTool workerFunctionRegistryTool;
     private final ConnectorManager connectorManager;
     private final TestingHttpServer server;
     private final CatalogManager catalogManager;
@@ -224,6 +228,11 @@ public TestingPrestoServer(List additionalModules)
         this(true, ImmutableMap.of(), null, null, new SqlParserOptions(), additionalModules);
     }
 
+    public TestingPrestoServer(Map properties) throws Exception
+    {
+        this(true, properties, null, null, new SqlParserOptions(), ImmutableList.of());
+    }
+
     public TestingPrestoServer(
             boolean coordinator,
             Map properties,
@@ -367,6 +376,9 @@ public TestingPrestoServer(
 
         connectorManager = injector.getInstance(ConnectorManager.class);
 
+        functionAndTypeManager = injector.getInstance(FunctionAndTypeManager.class);
+        workerFunctionRegistryTool = injector.getInstance(WorkerFunctionRegistryTool.class);
+
         server = injector.getInstance(TestingHttpServer.class);
         catalogManager = injector.getInstance(CatalogManager.class);
         transactionManager = injector.getInstance(TransactionManager.class);
@@ -501,6 +513,11 @@ private Map getServerProperties(
         return ImmutableMap.copyOf(serverProperties);
     }
 
+    public void registerWorkerFunctions()
+    {
+        functionAndTypeManager.registerWorkerFunctions(workerFunctionRegistryTool.getWorkerFunctions());
+    }
+
     @Override
     public void close()
             throws IOException
@@ -536,6 +553,12 @@ public void installCoordinatorPlugin(CoordinatorPlugin plugin)
         pluginManager.installCoordinatorPlugin(plugin);
     }
 
+    public void triggerConflictCheckWithBuiltInFunctions()
+    {
+        metadata.getFunctionAndTypeManager()
+                .getBuiltInPluginFunctionNamespaceManager().triggerConflictCheckWithBuiltInFunctions();
+    }
+
     public DispatchManager getDispatchManager()
     {
         return dispatchManager;
@@ -601,6 +624,11 @@ public HostAndPort getHttpsAddress()
         return HostAndPort.fromParts(httpsUri.getHost(), httpsUri.getPort());
     }
 
+    public URI getHttpBaseUrl()
+    {
+        return server.getHttpServerInfo().getHttpUri();
+    }
+
     public CatalogManager getCatalogManager()
     {
         return catalogManager;
diff --git a/presto-main/src/test/java/com/facebook/presto/TestClientRequestFilterPlugin.java b/presto-main/src/test/java/com/facebook/presto/TestClientRequestFilterPlugin.java
index d073f8ed96fee..eb295a80fbade 100644
--- a/presto-main/src/test/java/com/facebook/presto/TestClientRequestFilterPlugin.java
+++ b/presto-main/src/test/java/com/facebook/presto/TestClientRequestFilterPlugin.java
@@ -16,6 +16,7 @@
 import com.facebook.airlift.http.server.Authenticator;
 import com.facebook.presto.server.MockHttpServletRequest;
 import com.facebook.presto.server.security.AuthenticationFilter;
+import com.facebook.presto.server.security.DefaultWebUiAuthenticationManager;
 import com.facebook.presto.server.security.SecurityConfig;
 import com.facebook.presto.server.testing.TestingPrestoServer;
 import com.facebook.presto.spi.ClientRequestFilter;
@@ -111,7 +112,7 @@ private AuthenticationFilter setupAuthenticationFilter(List authenticators = createAuthenticators();
             SecurityConfig securityConfig = createSecurityConfig();
 
-            return new AuthenticationFilter(authenticators, securityConfig, clientRequestFilterManager);
+            return new AuthenticationFilter(authenticators, securityConfig, clientRequestFilterManager, new DefaultWebUiAuthenticationManager());
         }
     }
 
diff --git a/presto-main/src/test/java/com/facebook/presto/eventlistener/TestEventListenerManager.java b/presto-main/src/test/java/com/facebook/presto/eventlistener/TestEventListenerManager.java
index 9c21564fc41ef..4e6d4750f13da 100644
--- a/presto-main/src/test/java/com/facebook/presto/eventlistener/TestEventListenerManager.java
+++ b/presto-main/src/test/java/com/facebook/presto/eventlistener/TestEventListenerManager.java
@@ -25,6 +25,7 @@
 import com.facebook.presto.spi.eventlistener.EventListener;
 import com.facebook.presto.spi.eventlistener.EventListenerFactory;
 import com.facebook.presto.spi.eventlistener.OperatorStatistics;
+import com.facebook.presto.spi.eventlistener.OutputColumnMetadata;
 import com.facebook.presto.spi.eventlistener.PlanOptimizerInformation;
 import com.facebook.presto.spi.eventlistener.QueryCompletedEvent;
 import com.facebook.presto.spi.eventlistener.QueryContext;
@@ -374,10 +375,13 @@ private static QueryIOMetadata createDummyQueryIoMetadata()
         List inputs = new ArrayList<>();
         QueryInputMetadata queryInputMetadata = getQueryInputMetadata();
         inputs.add(queryInputMetadata);
-        Column column1 = new Column("column1", "int");
-        Column column2 = new Column("column2", "varchar");
-        Column column3 = new Column("column3", "varchar");
-        List columns = Arrays.asList(column1, column2, column3);
+        OutputColumnMetadata column1 = new OutputColumnMetadata("column1", "int", new HashSet<>());
+        OutputColumnMetadata column2 = new OutputColumnMetadata("column2", "varchar", new HashSet<>());
+        OutputColumnMetadata column3 = new OutputColumnMetadata("column3", "varchar", new HashSet<>());
+        List columns = new ArrayList<>();
+        columns.add(column1);
+        columns.add(column2);
+        columns.add(column3);
         QueryOutputMetadata outputMetadata = new QueryOutputMetadata(
                 "dummyCatalog",
                 "dummySchema",
diff --git a/presto-main/src/test/java/com/facebook/presto/server/MockHttpServletRequest.java b/presto-main/src/test/java/com/facebook/presto/server/MockHttpServletRequest.java
index b1fe03962f481..2033b62f88262 100644
--- a/presto-main/src/test/java/com/facebook/presto/server/MockHttpServletRequest.java
+++ b/presto-main/src/test/java/com/facebook/presto/server/MockHttpServletRequest.java
@@ -30,8 +30,10 @@
 import jakarta.servlet.http.HttpSession;
 import jakarta.servlet.http.HttpUpgradeHandler;
 import jakarta.servlet.http.Part;
+import jakarta.ws.rs.core.UriBuilder;
 
 import java.io.BufferedReader;
+import java.net.URI;
 import java.security.Principal;
 import java.util.Collection;
 import java.util.Enumeration;
@@ -50,12 +52,14 @@ public class MockHttpServletRequest
     private final ListMultimap headers;
     private final String remoteAddress;
     private final Map attributes;
+    private final String requestUrl;
 
     public MockHttpServletRequest(ListMultimap headers, String remoteAddress, Map attributes)
     {
         this.headers = ImmutableListMultimap.copyOf(requireNonNull(headers, "headers is null"));
         this.remoteAddress = requireNonNull(remoteAddress, "remoteAddress is null");
         this.attributes = new HashMap<>(requireNonNull(attributes, "attributes is null"));
+        this.requestUrl = null;
     }
 
     public MockHttpServletRequest(ListMultimap headers)
@@ -64,6 +68,14 @@ public MockHttpServletRequest(ListMultimap headers)
         this(headers, DEFAULT_ADDRESS, ImmutableMap.of());
     }
 
+    public MockHttpServletRequest(ListMultimap headers, String remoteAddress, String requestUrl)
+    {
+        this.headers = ImmutableListMultimap.copyOf(requireNonNull(headers, "headers is null"));
+        this.remoteAddress = requireNonNull(remoteAddress, "remoteAddress is null");
+        this.requestUrl = requireNonNull(requestUrl, "requestUrl is null");
+        this.attributes = ImmutableMap.of();
+    }
+
     @Override
     public String getAuthType()
     {
@@ -145,7 +157,12 @@ public String getContextPath()
     @Override
     public String getQueryString()
     {
-        throw new UnsupportedOperationException();
+        if (this.requestUrl == null) {
+            throw new UnsupportedOperationException();
+        }
+        URI uri = UriBuilder.fromUri(this.requestUrl).build();
+
+        return uri.getQuery();
     }
 
     @Override
@@ -181,7 +198,10 @@ public String getRequestURI()
     @Override
     public StringBuffer getRequestURL()
     {
-        throw new UnsupportedOperationException();
+        if (this.requestUrl == null) {
+            throw new UnsupportedOperationException();
+        }
+        return new StringBuffer(this.requestUrl);
     }
 
     @Override
@@ -337,7 +357,12 @@ public String getProtocol()
     @Override
     public String getScheme()
     {
-        throw new UnsupportedOperationException();
+        if (this.requestUrl == null) {
+            throw new UnsupportedOperationException();
+        }
+        URI uri = UriBuilder.fromUri(this.requestUrl).build();
+
+        return uri.getScheme();
     }
 
     @Override
diff --git a/presto-main/src/test/java/com/facebook/presto/server/security/oauth2/BaseOAuth2AuthenticationFilterTest.java b/presto-main/src/test/java/com/facebook/presto/server/security/oauth2/BaseOAuth2AuthenticationFilterTest.java
new file mode 100644
index 0000000000000..4540c655cb655
--- /dev/null
+++ b/presto-main/src/test/java/com/facebook/presto/server/security/oauth2/BaseOAuth2AuthenticationFilterTest.java
@@ -0,0 +1,403 @@
+/*
+ * Licensed 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 com.facebook.presto.server.security.oauth2;
+
+import com.facebook.airlift.json.JsonCodec;
+import com.facebook.airlift.log.Level;
+import com.facebook.airlift.log.Logging;
+import com.facebook.airlift.testing.Closeables;
+import com.facebook.presto.server.testing.TestingPrestoServer;
+import com.google.common.collect.ImmutableList;
+import com.google.common.collect.ImmutableMap;
+import com.google.inject.Key;
+import io.jsonwebtoken.impl.DefaultClaims;
+import okhttp3.Cookie;
+import okhttp3.CookieJar;
+import okhttp3.HttpUrl;
+import okhttp3.JavaNetCookieJar;
+import okhttp3.OkHttpClient;
+import okhttp3.Request;
+import okhttp3.Response;
+import org.testng.annotations.AfterClass;
+import org.testng.annotations.BeforeClass;
+import org.testng.annotations.Test;
+
+import java.io.IOException;
+import java.net.CookieManager;
+import java.net.CookieStore;
+import java.net.HttpCookie;
+import java.net.MalformedURLException;
+import java.net.URI;
+import java.net.URL;
+import java.security.KeyPairGenerator;
+import java.security.NoSuchAlgorithmException;
+import java.time.Duration;
+import java.time.Instant;
+import java.util.List;
+import java.util.Optional;
+import java.util.UUID;
+
+import static com.facebook.airlift.testing.Assertions.assertLessThan;
+import static com.facebook.airlift.units.Duration.nanosSince;
+import static com.facebook.presto.client.OkHttpUtil.setupInsecureSsl;
+import static com.facebook.presto.server.security.oauth2.JwtUtil.newJwtBuilder;
+import static com.facebook.presto.server.security.oauth2.OAuthWebUiCookie.OAUTH2_COOKIE;
+import static com.facebook.presto.server.security.oauth2.TokenEndpointAuthMethod.CLIENT_SECRET_BASIC;
+import static jakarta.servlet.http.HttpServletResponse.SC_OK;
+import static jakarta.ws.rs.core.HttpHeaders.AUTHORIZATION;
+import static jakarta.ws.rs.core.HttpHeaders.LOCATION;
+import static jakarta.ws.rs.core.Response.Status.OK;
+import static jakarta.ws.rs.core.Response.Status.SEE_OTHER;
+import static jakarta.ws.rs.core.Response.Status.UNAUTHORIZED;
+import static java.util.concurrent.TimeUnit.MILLISECONDS;
+import static java.util.concurrent.TimeUnit.SECONDS;
+import static org.assertj.core.api.Assertions.assertThat;
+import static org.assertj.core.api.Assertions.fail;
+import static org.testng.Assert.assertEquals;
+
+public abstract class BaseOAuth2AuthenticationFilterTest
+{
+    protected static final Duration TTL_ACCESS_TOKEN_IN_SECONDS = Duration.ofSeconds(5);
+
+    protected static final String PRESTO_CLIENT_ID = "presto-client";
+    protected static final String PRESTO_CLIENT_SECRET = "presto-secret";
+    private static final String PRESTO_AUDIENCE = PRESTO_CLIENT_ID;
+    private static final String ADDITIONAL_AUDIENCE = "https://external-service.com";
+    protected static final String TRUSTED_CLIENT_ID = "trusted-client";
+    protected static final String TRUSTED_CLIENT_SECRET = "trusted-secret";
+    private static final String UNTRUSTED_CLIENT_ID = "untrusted-client";
+    private static final String UNTRUSTED_CLIENT_SECRET = "untrusted-secret";
+    private static final String UNTRUSTED_CLIENT_AUDIENCE = "https://untrusted.com";
+
+    private final Logging logging = Logging.initialize();
+    protected final OkHttpClient httpClient;
+    protected TestingHydraIdentityProvider hydraIdP;
+
+    private TestingPrestoServer server;
+
+    private SimpleProxyServer simpleProxy;
+    private URI uiUri;
+
+    private URI proxyURI;
+
+    protected BaseOAuth2AuthenticationFilterTest()
+    {
+        OkHttpClient.Builder httpClientBuilder = new OkHttpClient.Builder();
+        setupInsecureSsl(httpClientBuilder);
+        httpClientBuilder.followRedirects(false);
+        httpClient = httpClientBuilder.build();
+    }
+
+    static void waitForNodeRefresh(TestingPrestoServer server)
+            throws InterruptedException
+    {
+        long start = System.nanoTime();
+        while (server.refreshNodes().getActiveNodes().size() < 1) {
+            assertLessThan(nanosSince(start), new com.facebook.airlift.units.Duration(10, SECONDS));
+            MILLISECONDS.sleep(10);
+        }
+    }
+
+    @BeforeClass
+    public void setup()
+            throws Exception
+    {
+        logging.setLevel(OAuth2Service.class.getName(), Level.DEBUG);
+
+        hydraIdP = getHydraIdp();
+        String idpUrl = "https://localhost:" + hydraIdP.getAuthPort();
+        server = new TestingPrestoServer(getOAuth2Config(idpUrl));
+        server.getInstance(Key.get(OAuth2Client.class)).load();
+        waitForNodeRefresh(server);
+        // Due to problems with the Presto OSS project related to the AuthenticationFilter we have to run Presto behind a Proxy and terminate SSL at the proxy.
+        simpleProxy = new SimpleProxyServer(server.getHttpBaseUrl());
+        MILLISECONDS.sleep(1000);
+        proxyURI = URI.create("https://127.0.0.1:" + simpleProxy.getHttpsBaseUrl().getPort());
+        uiUri = proxyURI.resolve("/");
+
+        hydraIdP.createClient(
+                PRESTO_CLIENT_ID,
+                PRESTO_CLIENT_SECRET,
+                CLIENT_SECRET_BASIC,
+                ImmutableList.of(PRESTO_AUDIENCE, ADDITIONAL_AUDIENCE),
+                proxyURI + "/oauth2/callback");
+        hydraIdP.createClient(
+                TRUSTED_CLIENT_ID,
+                TRUSTED_CLIENT_SECRET,
+                CLIENT_SECRET_BASIC,
+                ImmutableList.of(TRUSTED_CLIENT_ID),
+                proxyURI + "/oauth2/callback");
+        hydraIdP.createClient(
+                UNTRUSTED_CLIENT_ID,
+                UNTRUSTED_CLIENT_SECRET,
+                CLIENT_SECRET_BASIC,
+                ImmutableList.of(UNTRUSTED_CLIENT_AUDIENCE),
+                "https://untrusted.com/callback");
+    }
+
+    protected abstract ImmutableMap getOAuth2Config(String idpUrl)
+            throws IOException;
+
+    protected abstract TestingHydraIdentityProvider getHydraIdp()
+            throws Exception;
+
+    @AfterClass(alwaysRun = true)
+    public void tearDown()
+            throws Exception
+    {
+        Closeables.closeAll(server, hydraIdP, simpleProxy);
+    }
+
+    @Test
+    public void testUnauthorizedApiCall()
+            throws IOException
+    {
+        try (Response response = httpClient
+                .newCall(apiCall().build())
+                .execute()) {
+            assertUnauthorizedResponse(response);
+        }
+    }
+
+    @Test
+    public void testUnauthorizedUICall()
+            throws IOException
+    {
+        try (Response response = httpClient
+                .newCall(uiCall().build())
+                .execute()) {
+            assertRedirectResponse(response);
+        }
+    }
+
+    @Test
+    public void testUnsignedToken()
+            throws NoSuchAlgorithmException, IOException
+    {
+        KeyPairGenerator keyGenerator = KeyPairGenerator.getInstance("RSA");
+        keyGenerator.initialize(4096);
+        long now = Instant.now().getEpochSecond();
+        String token = newJwtBuilder()
+                .setHeaderParam("alg", "RS256")
+                .setHeaderParam("kid", "public:f467aa08-1c1b-4cde-ba45-84b0ef5d2ba8")
+                .setHeaderParam("typ", "JWT")
+                .setClaims(
+                        new DefaultClaims(
+                                ImmutableMap.builder()
+                                        .put("aud", ImmutableList.of())
+                                        .put("client_id", PRESTO_CLIENT_ID)
+                                        .put("exp", now + 60L)
+                                        .put("iat", now)
+                                        .put("iss", "https://hydra:4444/")
+                                        .put("jti", UUID.randomUUID())
+                                        .put("nbf", now)
+                                        .put("scp", ImmutableList.of("openid"))
+                                        .put("sub", "foo@bar.com")
+                                        .build()))
+                .signWith(keyGenerator.generateKeyPair().getPrivate())
+                .compact();
+        try (Response response = httpClientWithOAuth2Cookie(token, false)
+                .newCall(apiCall().build())
+                .execute()) {
+            assertUnauthorizedResponse(response);
+        }
+    }
+
+    @Test
+    public void testTokenWithInvalidAudience()
+            throws IOException
+    {
+        String token = hydraIdP.getToken(UNTRUSTED_CLIENT_ID, UNTRUSTED_CLIENT_SECRET, ImmutableList.of(UNTRUSTED_CLIENT_AUDIENCE));
+        try (Response response = httpClientWithOAuth2Cookie(token, false)
+                .newCall(apiCall().build())
+                .execute()) {
+            assertUnauthorizedResponse(response);
+        }
+    }
+
+    @Test
+    public void testTokenFromTrustedClient()
+            throws IOException
+    {
+        String token = hydraIdP.getToken(TRUSTED_CLIENT_ID, TRUSTED_CLIENT_SECRET, ImmutableList.of(TRUSTED_CLIENT_ID));
+        assertUICallWithCookie(token);
+    }
+
+    @Test
+    public void testTokenWithMultipleAudiences()
+            throws IOException
+    {
+        String token = hydraIdP.getToken(PRESTO_CLIENT_ID, PRESTO_CLIENT_SECRET, ImmutableList.of(PRESTO_AUDIENCE, ADDITIONAL_AUDIENCE));
+        assertUICallWithCookie(token);
+    }
+
+    @Test
+    public void testSuccessfulFlow()
+            throws Exception
+    {
+        // create a new HttpClient which follows redirects and give access to cookies
+        CookieManager cookieManager = new CookieManager();
+        CookieStore cookieStore = cookieManager.getCookieStore();
+        OkHttpClient.Builder httpClientBuilder = new OkHttpClient.Builder();
+        setupInsecureSsl(httpClientBuilder);
+        OkHttpClient httpClient = httpClientBuilder
+                .followRedirects(true)
+                .cookieJar(new JavaNetCookieJar(cookieManager))
+                .build();
+
+        assertThat(cookieStore.get(uiUri)).isEmpty();
+
+        // access UI and follow redirects in order to get OAuth2 cookie
+        Response response = httpClient.newCall(
+                new Request.Builder()
+                        .url(uiUri.toURL())
+                        .get()
+                        .build())
+                .execute();
+
+        assertEquals(response.code(), SC_OK);
+
+        Optional oauth2Cookie = cookieStore.get(uiUri)
+                .stream()
+                .filter(cookie -> cookie.getName().equals(OAUTH2_COOKIE))
+                .findFirst();
+        assertThat(oauth2Cookie).isNotEmpty();
+        assertOAuth2Cookie(oauth2Cookie.get());
+        assertUICallWithCookie(oauth2Cookie.get().getValue());
+    }
+
+    @Test
+    public void testExpiredAccessToken()
+            throws Exception
+    {
+        String token = hydraIdP.getToken(PRESTO_CLIENT_ID, PRESTO_CLIENT_SECRET, ImmutableList.of(PRESTO_AUDIENCE));
+        assertUICallWithCookie(token);
+        Thread.sleep(TTL_ACCESS_TOKEN_IN_SECONDS.plusSeconds(1).toMillis()); // wait for the token expiration = ttl of access token + 1 sec
+        try (Response response = httpClientWithOAuth2Cookie(token, false).newCall(apiCall().build()).execute()) {
+            assertUnauthorizedResponse(response);
+        }
+    }
+
+    private Request.Builder uiCall()
+    {
+        return new Request.Builder()
+                .url(proxyURI.resolve("/").toString())
+                .get();
+    }
+
+    private Request.Builder apiCall()
+    {
+        return new Request.Builder()
+                .url(proxyURI.resolve("/v1/cluster").toString())
+                .get();
+    }
+
+    private void assertOAuth2Cookie(HttpCookie cookie)
+    {
+        assertThat(cookie.getName()).isEqualTo(OAUTH2_COOKIE);
+        assertThat(cookie.getDomain()).isIn(proxyURI.getHost());
+        assertThat(cookie.getPath()).isEqualTo("/");
+        assertThat(cookie.getSecure()).isTrue();
+        assertThat(cookie.isHttpOnly()).isTrue();
+        assertThat(cookie.getMaxAge()).isLessThanOrEqualTo(TTL_ACCESS_TOKEN_IN_SECONDS.getSeconds());
+        validateAccessToken(cookie.getValue());
+    }
+
+    protected void validateAccessToken(String cookieValue)
+    {
+        Request request = new Request.Builder().url("https://localhost:" + hydraIdP.getAuthPort() + "/userinfo").addHeader(AUTHORIZATION, "Bearer " + cookieValue).build();
+        try (Response response = httpClient.newCall(request).execute()) {
+            assertThat(response.body()).isNotNull();
+            DefaultClaims claims = new DefaultClaims(JsonCodec.mapJsonCodec(String.class, Object.class).fromJson(response.body().bytes()));
+            assertThat(claims.getSubject()).isEqualTo("foo@bar.com");
+        }
+        catch (IOException e) {
+            fail("Exception while calling /userinfo", e);
+        }
+    }
+
+    private void assertUICallWithCookie(String cookieValue)
+            throws IOException
+    {
+        OkHttpClient httpClient = httpClientWithOAuth2Cookie(cookieValue, true);
+        // pass access token in Presto UI cookie
+        try (Response response = httpClient.newCall(uiCall().build())
+                .execute()) {
+            assertThat(response.code()).isEqualTo(OK.getStatusCode());
+        }
+    }
+
+    @SuppressWarnings("NullableProblems")
+    private OkHttpClient httpClientWithOAuth2Cookie(String cookieValue, boolean followRedirects)
+    {
+        OkHttpClient.Builder httpClientBuilder = new OkHttpClient.Builder();
+        setupInsecureSsl(httpClientBuilder);
+        httpClientBuilder.followRedirects(followRedirects);
+        httpClientBuilder.cookieJar(new CookieJar()
+        {
+            @Override
+            public void saveFromResponse(HttpUrl url, List cookies)
+            {
+            }
+
+            @Override
+            public List loadForRequest(HttpUrl url)
+            {
+                Cookie cookie = new Cookie.Builder()
+                            .domain(proxyURI.getHost())
+                            .path("/")
+                            .name(OAUTH2_COOKIE)
+                            .value(cookieValue)
+                            .secure()
+                            .build();
+                return ImmutableList.of(cookie);
+            }
+        });
+        return httpClientBuilder.build();
+    }
+
+    private void assertRedirectResponse(Response response)
+            throws MalformedURLException
+    {
+        assertThat(response.code()).isEqualTo(SEE_OTHER.getStatusCode());
+        assertRedirectUrl(response.header(LOCATION));
+    }
+
+    private void assertUnauthorizedResponse(Response response)
+            throws IOException
+    {
+        assertThat(response.code()).isEqualTo(UNAUTHORIZED.getStatusCode());
+        assertThat(response.body()).isNotNull();
+        // NOTE that our errors come in looking like an HTML page since we don't do anything special on the server side so it just is like that.
+        assertThat(response.body().string()).contains("Invalid Credentials");
+    }
+
+    private void assertRedirectUrl(String redirectUrl)
+            throws MalformedURLException
+    {
+        assertThat(redirectUrl).isNotNull();
+        URL location = new URL(redirectUrl);
+        HttpUrl url = HttpUrl.parse(redirectUrl);
+        assertThat(url).isNotNull();
+        assertThat(location.getProtocol()).isEqualTo("https");
+        assertThat(location.getHost()).isEqualTo("localhost");
+        assertThat(location.getPort()).isEqualTo(hydraIdP.getAuthPort());
+        assertThat(location.getPath()).isEqualTo("/oauth2/auth");
+        assertThat(url.queryParameterValues("response_type")).isEqualTo(ImmutableList.of("code"));
+        assertThat(url.queryParameterValues("scope")).isEqualTo(ImmutableList.of("openid"));
+        assertThat(url.queryParameterValues("redirect_uri")).isEqualTo(ImmutableList.of(proxyURI + "/oauth2/callback"));
+        assertThat(url.queryParameterValues("client_id")).isEqualTo(ImmutableList.of(PRESTO_CLIENT_ID));
+        assertThat(url.queryParameterValues("state")).isNotNull();
+    }
+}
diff --git a/presto-main/src/test/java/com/facebook/presto/server/security/oauth2/SimpleProxyServer.java b/presto-main/src/test/java/com/facebook/presto/server/security/oauth2/SimpleProxyServer.java
new file mode 100644
index 0000000000000..99e8bb43c4b8d
--- /dev/null
+++ b/presto-main/src/test/java/com/facebook/presto/server/security/oauth2/SimpleProxyServer.java
@@ -0,0 +1,196 @@
+/*
+ * Licensed 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 com.facebook.presto.server.security.oauth2;
+
+import com.facebook.airlift.http.server.HttpServerConfig;
+import com.facebook.airlift.http.server.HttpServerInfo;
+import com.facebook.airlift.http.server.testing.TestingHttpServer;
+import com.facebook.airlift.log.Logger;
+import com.facebook.airlift.node.NodeInfo;
+import com.google.common.collect.ImmutableMap;
+import com.google.common.io.Resources;
+import jakarta.servlet.ServletException;
+import jakarta.servlet.http.Cookie;
+import jakarta.servlet.http.HttpServlet;
+import jakarta.servlet.http.HttpServletRequest;
+import jakarta.servlet.http.HttpServletResponse;
+import jakarta.ws.rs.core.UriBuilder;
+import okhttp3.Headers;
+import okhttp3.OkHttpClient;
+import okhttp3.Request;
+import okhttp3.Response;
+
+import java.io.Closeable;
+import java.io.IOException;
+import java.io.InputStream;
+import java.io.OutputStream;
+import java.net.URI;
+import java.util.List;
+import java.util.Optional;
+
+import static com.facebook.presto.client.OkHttpUtil.setupInsecureSsl;
+import static com.facebook.presto.server.security.oauth2.OAuth2Utils.getFullRequestURL;
+import static com.google.common.base.Throwables.throwIfUnchecked;
+import static com.google.common.net.HttpHeaders.HOST;
+import static com.google.common.net.HttpHeaders.X_FORWARDED_FOR;
+import static com.google.common.net.HttpHeaders.X_FORWARDED_PROTO;
+import static java.lang.String.format;
+import static java.util.concurrent.TimeUnit.SECONDS;
+
+public class SimpleProxyServer
+        implements Closeable
+{
+    private final TestingHttpServer server;
+
+    public SimpleProxyServer(URI forwardBaseURI)
+            throws Exception
+    {
+        server = createSimpleProxyServer(forwardBaseURI);
+        server.start();
+    }
+
+    @Override
+    public void close()
+            throws IOException
+    {
+        try {
+            server.stop();
+        }
+        catch (Exception e) {
+            throwIfUnchecked(e);
+            throw new RuntimeException(e);
+        }
+    }
+
+    public URI getHttpsBaseUrl()
+    {
+        return server.getHttpServerInfo().getHttpsUri();
+    }
+
+    private TestingHttpServer createSimpleProxyServer(URI forwardBaseURI)
+            throws IOException
+    {
+        NodeInfo nodeInfo = new NodeInfo("test");
+        HttpServerConfig config = new HttpServerConfig()
+                .setHttpPort(0)
+                .setHttpsEnabled(true)
+                .setHttpsPort(0)
+                .setKeystorePath(Resources.getResource("cert/localhost.pem").getPath());
+        HttpServerInfo httpServerInfo = new HttpServerInfo(config, nodeInfo);
+        return new TestingHttpServer(httpServerInfo, nodeInfo, config, new SimpleProxy(forwardBaseURI), ImmutableMap.of(), ImmutableMap.of(), Optional.empty());
+    }
+
+    private class SimpleProxy
+            extends HttpServlet
+    {
+        private final OkHttpClient httpClient;
+        private final URI forwardBaseURI;
+
+        private final Logger logger = Logger.get(SimpleProxy.class);
+
+        public SimpleProxy(URI forwardBaseURI)
+        {
+            this.forwardBaseURI = forwardBaseURI;
+            OkHttpClient.Builder httpClientBuilder = new OkHttpClient.Builder();
+            setupInsecureSsl(httpClientBuilder);
+            httpClient = httpClientBuilder
+                    .followRedirects(false)
+                    .connectTimeout(100, SECONDS)
+                    .writeTimeout(100, SECONDS)
+                    .readTimeout(100, SECONDS)
+                    .build();
+        }
+
+        @Override
+        protected void service(HttpServletRequest request, HttpServletResponse servletResponse)
+                throws ServletException, IOException
+        {
+            UriBuilder requestUriBuilder = UriBuilder.fromUri(getFullRequestURL(request));
+            requestUriBuilder
+                    .scheme("http")
+                    .host(forwardBaseURI.getHost())
+                    .port(forwardBaseURI.getPort());
+
+            String hostHeader = new StringBuilder().append(request.getRemoteHost()).append(":").append(request.getLocalPort()).toString();
+            Cookie[] cookies = Optional.ofNullable(request.getCookies()).orElse(new Cookie[0]);
+            String requestUri = requestUriBuilder.build().toString();
+            Request.Builder reqBuilder = new Request.Builder()
+                    .url(requestUri)
+                    .addHeader(X_FORWARDED_PROTO, "https")
+                    .addHeader(X_FORWARDED_FOR, request.getRemoteAddr())
+                    .addHeader(HOST, hostHeader)
+                    .get();
+
+            if (cookies.length > 0) {
+                for (Cookie cookie : cookies) {
+                    reqBuilder.addHeader("Cookie", cookie.getName() + "=" + cookie.getValue());
+                }
+            }
+            Response response;
+            try {
+                response = httpClient.newCall(reqBuilder.build()).execute();
+                servletResponse.setStatus(response.code());
+
+                Headers responseHeaders = response.headers();
+                responseHeaders.names().stream().forEach(headerName -> {
+                    // Headers can have multiple values
+                    List headerValues = responseHeaders.values(headerName);
+                    headerValues.forEach(headerValue -> {
+                        servletResponse.addHeader(headerName, headerValue);
+                    });
+                });
+
+                //copy the response body to the servlet response.
+                InputStream is = response.body().byteStream();
+                OutputStream os = servletResponse.getOutputStream();
+                byte[] buffer = new byte[10 * 1024];
+                int read;
+                while ((read = is.read(buffer)) != -1) {
+                    os.write(buffer, 0, read);
+                }
+            }
+            catch (Exception e) {
+                logger.error(format("Encountered an error while proxying request to %s", requestUri), e);
+                servletResponse.sendError(HttpServletResponse.SC_INTERNAL_SERVER_ERROR);
+            }
+        }
+    }
+
+    // This is just to help iterate changes that might need to be made in the future to the Simple Proxy Server for test purposes.
+    // Waiting for the tests to run can be really slow so having this helper here is nice if you want quicker feedback.
+    private static void runTestServer()
+            throws Exception
+    {
+        SimpleProxyServer test = new SimpleProxyServer(new URI(""));
+        OkHttpClient.Builder httpClientBuilder = new OkHttpClient.Builder();
+        setupInsecureSsl(httpClientBuilder);
+        OkHttpClient client = httpClientBuilder.build();
+        Request.Builder req = new Request.Builder().url(test.getHttpsBaseUrl().resolve("/v1/query").toString()).get();
+        Logger logger = Logger.get("Run Test Server Debug Helper");
+        try {
+            client.newCall(req.build()).execute();
+        }
+        catch (Exception e) {
+            logger.error(e);
+        }
+
+        test.close();
+    }
+
+    public static void main(String[] args)
+            throws Exception
+    {
+        runTestServer();
+    }
+}
diff --git a/presto-main/src/test/java/com/facebook/presto/server/security/oauth2/TestDualAuthenticationFilterWithOAuth.java b/presto-main/src/test/java/com/facebook/presto/server/security/oauth2/TestDualAuthenticationFilterWithOAuth.java
new file mode 100644
index 0000000000000..089c6b9c80aaf
--- /dev/null
+++ b/presto-main/src/test/java/com/facebook/presto/server/security/oauth2/TestDualAuthenticationFilterWithOAuth.java
@@ -0,0 +1,249 @@
+/*
+ * Licensed 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 com.facebook.presto.server.security.oauth2;
+
+import com.facebook.airlift.log.Level;
+import com.facebook.airlift.log.Logging;
+import com.facebook.presto.server.testing.TestingPrestoServer;
+import com.google.common.collect.ImmutableList;
+import com.google.common.collect.ImmutableMap;
+import com.google.common.io.Resources;
+import com.google.inject.Key;
+import okhttp3.Cookie;
+import okhttp3.CookieJar;
+import okhttp3.HttpUrl;
+import okhttp3.OkHttpClient;
+import okhttp3.Request;
+import okhttp3.Response;
+import org.testng.annotations.BeforeClass;
+import org.testng.annotations.Test;
+
+import java.io.IOException;
+import java.net.URI;
+import java.time.Duration;
+import java.util.Base64;
+import java.util.List;
+
+import static com.facebook.airlift.testing.Assertions.assertLessThan;
+import static com.facebook.airlift.units.Duration.nanosSince;
+import static com.facebook.presto.client.OkHttpUtil.setupInsecureSsl;
+import static com.facebook.presto.server.security.oauth2.OAuthWebUiCookie.OAUTH2_COOKIE;
+import static com.facebook.presto.server.security.oauth2.TokenEndpointAuthMethod.CLIENT_SECRET_BASIC;
+import static com.google.common.net.HttpHeaders.AUTHORIZATION;
+import static com.google.common.net.HttpHeaders.WWW_AUTHENTICATE;
+import static jakarta.ws.rs.core.Response.Status.OK;
+import static jakarta.ws.rs.core.Response.Status.UNAUTHORIZED;
+import static java.io.File.createTempFile;
+import static java.util.concurrent.TimeUnit.MILLISECONDS;
+import static java.util.concurrent.TimeUnit.SECONDS;
+import static org.assertj.core.api.Assertions.assertThat;
+
+public class TestDualAuthenticationFilterWithOAuth
+{
+    protected static final Duration TTL_ACCESS_TOKEN_IN_SECONDS = Duration.ofSeconds(5);
+    protected static final String PRESTO_CLIENT_ID = "presto-client";
+    protected static final String PRESTO_CLIENT_SECRET = "presto-secret";
+    private static final String PRESTO_AUDIENCE = PRESTO_CLIENT_ID;
+    private static final String ADDITIONAL_AUDIENCE = "https://external-service.com";
+    protected static final String TRUSTED_CLIENT_ID = "trusted-client";
+
+    private final Logging logging = Logging.initialize();
+    protected final OkHttpClient httpClient;
+    protected TestingHydraIdentityProvider hydraIdP;
+
+    private TestingPrestoServer server;
+
+    private SimpleProxyServer simpleProxy;
+
+    private URI proxyURI;
+
+    protected TestDualAuthenticationFilterWithOAuth()
+    {
+        OkHttpClient.Builder httpClientBuilder = new OkHttpClient.Builder();
+        setupInsecureSsl(httpClientBuilder);
+        httpClientBuilder.followRedirects(false);
+        httpClient = httpClientBuilder.build();
+    }
+
+    static void waitForNodeRefresh(TestingPrestoServer server)
+            throws InterruptedException
+    {
+        long start = System.nanoTime();
+        while (server.refreshNodes().getActiveNodes().size() < 1) {
+            assertLessThan(nanosSince(start), new com.facebook.airlift.units.Duration(10, SECONDS));
+            MILLISECONDS.sleep(10);
+        }
+    }
+
+    protected ImmutableMap getConfig(String idpUrl)
+            throws IOException
+    {
+        return ImmutableMap.builder()
+                .put("http-server.authentication.allow-forwarded-https", "true")
+                .put("http-server.authentication.type", "OAUTH2,PASSWORD")
+                .put("http-server.authentication.oauth2.issuer", "https://localhost:4444/")
+                .put("http-server.authentication.oauth2.auth-url", idpUrl + "/oauth2/auth")
+                .put("http-server.authentication.oauth2.token-url", idpUrl + "/oauth2/token")
+                .put("http-server.authentication.oauth2.jwks-url", idpUrl + "/.well-known/jwks.json")
+                .put("http-server.authentication.oauth2.client-id", PRESTO_CLIENT_ID)
+                .put("http-server.authentication.oauth2.client-secret", PRESTO_CLIENT_SECRET)
+                .put("http-server.authentication.oauth2.additional-audiences", TRUSTED_CLIENT_ID)
+                .put("http-server.authentication.oauth2.max-clock-skew", "0s")
+                .put("http-server.authentication.oauth2.user-mapping.pattern", "(.*)(@.*)?")
+                .put("http-server.authentication.oauth2.oidc.discovery", "false")
+                .put("configuration-based-authorizer.role-regex-map.file-path", createTempFile("regex-map", null).getAbsolutePath().toString())
+                .put("oauth2-jwk.http-client.trust-store-path", Resources.getResource("cert/localhost.pem").getPath())
+                .build();
+    }
+
+    @BeforeClass
+    public void setup()
+            throws Exception
+    {
+        logging.setLevel(OAuth2Service.class.getName(), Level.DEBUG);
+        hydraIdP = new TestingHydraIdentityProvider(TTL_ACCESS_TOKEN_IN_SECONDS, true, false);
+        hydraIdP.start();
+        String idpUrl = "https://localhost:" + hydraIdP.getAuthPort();
+        server = new TestingPrestoServer(getConfig(idpUrl));
+        server.getInstance(Key.get(OAuth2Client.class)).load();
+        waitForNodeRefresh(server);
+        // Due to problems with the Presto OSS project related to the AuthenticationFilter we have to run Presto behind a Proxy and terminate SSL at the proxy.
+        simpleProxy = new SimpleProxyServer(server.getHttpBaseUrl());
+        MILLISECONDS.sleep(1000);
+        proxyURI = URI.create("https://127.0.0.1:" + simpleProxy.getHttpsBaseUrl().getPort());
+
+        hydraIdP.createClient(
+                PRESTO_CLIENT_ID,
+                PRESTO_CLIENT_SECRET,
+                CLIENT_SECRET_BASIC,
+                ImmutableList.of(PRESTO_AUDIENCE, ADDITIONAL_AUDIENCE),
+                simpleProxy.getHttpsBaseUrl() + "/oauth2/callback");
+    }
+
+    @Test
+    public void testExpiredOAuthToken()
+            throws Exception
+    {
+        String token = hydraIdP.getToken(PRESTO_CLIENT_ID, PRESTO_CLIENT_SECRET, ImmutableList.of(PRESTO_AUDIENCE));
+        assertUICallWithCookie(token);
+        Thread.sleep(TTL_ACCESS_TOKEN_IN_SECONDS.plusSeconds(1).toMillis()); // wait for the token expiration = ttl of access token + 1 sec
+        try (Response response = httpClientWithOAuth2Cookie(token, false).newCall(apiCall().build()).execute()) {
+            assertUnauthorizedOAuthOnlyHeaders(response);
+        }
+    }
+
+    @Test
+    public void testNoAuth()
+            throws Exception
+    {
+        try (Response response = httpClient
+                .newCall(apiCall().build())
+                .execute()) {
+            assertAllUnauthorizedHeaders(response);
+        }
+    }
+
+    @Test
+    public void testInvalidBasicAuth()
+            throws Exception
+    {
+        String userPass = "test:password";
+        String basicAuth = "Basic " + Base64.getEncoder().encodeToString(userPass.getBytes());
+        try (Response response = httpClient
+                .newCall(apiCall().addHeader(AUTHORIZATION, basicAuth).build())
+                .execute()) {
+            assertAllUnauthorizedHeaders(response);
+        }
+    }
+
+    private Request.Builder apiCall()
+    {
+        return new Request.Builder()
+                .url(proxyURI.resolve("/v1/cluster").toString())
+                .get();
+    }
+
+    private void assertUnauthorizedOAuthOnlyHeaders(Response response)
+            throws IOException
+    {
+        String redirectServer = "x_redirect_server=\"" + proxyURI.resolve("/oauth2/token/initiate/");
+        String tokenServer = "x_token_server=\"" + proxyURI.resolve("/oauth2/token/");
+        assertUnauthorizedResponse(response);
+        List headers = response.headers(WWW_AUTHENTICATE);
+        assertThat(headers.size()).isEqualTo(1);
+        assertThat(headers.get(0)).contains(tokenServer, redirectServer);
+    }
+
+    private void assertAllUnauthorizedHeaders(Response response)
+            throws IOException
+    {
+        String redirectServer = "x_redirect_server=\"" + proxyURI.resolve("/oauth2/token/initiate/").toString();
+        String tokenServer = "x_token_server=\"" + proxyURI.resolve("/oauth2/token/");
+        assertUnauthorizedResponse(response);
+        List headers = response.headers(WWW_AUTHENTICATE);
+        assertThat(headers.size()).isEqualTo(2);
+        assertThat(headers.stream().allMatch(h ->
+                        h.contains("Basic realm=\"Presto\"") ||
+                                (h.contains(redirectServer) && h.contains(tokenServer))
+                )).isTrue();
+    }
+
+    private void assertUICallWithCookie(String cookieValue)
+            throws IOException
+    {
+        OkHttpClient httpClient = httpClientWithOAuth2Cookie(cookieValue, true);
+        // pass access token in Presto UI cookie
+        try (Response response = httpClient.newCall(apiCall().build())
+                .execute()) {
+            assertThat(response.code()).isEqualTo(OK.getStatusCode());
+        }
+    }
+
+    private OkHttpClient httpClientWithOAuth2Cookie(String cookieValue, boolean followRedirects)
+    {
+        OkHttpClient.Builder httpClientBuilder = new OkHttpClient.Builder();
+        setupInsecureSsl(httpClientBuilder);
+        httpClientBuilder.followRedirects(followRedirects);
+        httpClientBuilder.cookieJar(new CookieJar()
+        {
+            @Override
+            public void saveFromResponse(HttpUrl url, List cookies)
+            {
+            }
+
+            @Override
+            public List loadForRequest(HttpUrl url)
+            {
+                return ImmutableList.of(new Cookie.Builder()
+                        .domain(proxyURI.getHost())
+                        .path("/")
+                        .name(OAUTH2_COOKIE)
+                        .value(cookieValue)
+                        .secure()
+                        .build());
+            }
+        });
+        return httpClientBuilder.build();
+    }
+
+    private void assertUnauthorizedResponse(Response response)
+            throws IOException
+    {
+        assertThat(response.code()).isEqualTo(UNAUTHORIZED.getStatusCode());
+        assertThat(response.body()).isNotNull();
+        // NOTE that our errors come in looking like an HTML page since we don't do anything special on the server side so it just is like that.
+        assertThat(response.body().string()).contains("Invalid Credentials");
+    }
+}
diff --git a/presto-main/src/test/java/com/facebook/presto/server/security/oauth2/TestJweTokenSerializer.java b/presto-main/src/test/java/com/facebook/presto/server/security/oauth2/TestJweTokenSerializer.java
new file mode 100644
index 0000000000000..7a41c0fdaeeea
--- /dev/null
+++ b/presto-main/src/test/java/com/facebook/presto/server/security/oauth2/TestJweTokenSerializer.java
@@ -0,0 +1,171 @@
+/*
+ * Licensed 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 com.facebook.presto.server.security.oauth2;
+
+import com.facebook.airlift.units.Duration;
+import com.facebook.presto.server.security.oauth2.TokenPairSerializer.TokenPair;
+import com.nimbusds.jose.KeyLengthException;
+import io.jsonwebtoken.ExpiredJwtException;
+import io.jsonwebtoken.Jwts;
+import org.testng.annotations.Test;
+
+import java.net.URI;
+import java.security.GeneralSecurityException;
+import java.time.Clock;
+import java.time.Instant;
+import java.time.ZoneId;
+import java.time.ZonedDateTime;
+import java.util.Calendar;
+import java.util.Date;
+import java.util.Map;
+import java.util.Optional;
+
+import static com.facebook.airlift.units.Duration.succinctDuration;
+import static com.facebook.presto.server.security.oauth2.TokenPairSerializer.TokenPair.accessAndRefreshTokens;
+import static java.time.temporal.ChronoUnit.MILLIS;
+import static java.util.concurrent.TimeUnit.MINUTES;
+import static java.util.concurrent.TimeUnit.SECONDS;
+import static org.assertj.core.api.Assertions.assertThat;
+import static org.assertj.core.api.Assertions.assertThatThrownBy;
+
+public class TestJweTokenSerializer
+{
+    @Test
+    public void testSerialization()
+            throws Exception
+    {
+        JweTokenSerializer serializer = tokenSerializer(Clock.systemUTC(), succinctDuration(5, SECONDS));
+
+        Date expiration = new Calendar.Builder().setDate(2022, 6, 22).build().getTime();
+        String serializedTokenPair = serializer.serialize(accessAndRefreshTokens("access_token", expiration, "refresh_token"));
+        TokenPair deserializedTokenPair = serializer.deserialize(serializedTokenPair);
+
+        assertThat(deserializedTokenPair.getAccessToken()).isEqualTo("access_token");
+        assertThat(deserializedTokenPair.getExpiration()).isEqualTo(expiration);
+        assertThat(deserializedTokenPair.getRefreshToken()).isEqualTo(Optional.of("refresh_token"));
+    }
+
+    @Test
+    public void testTokenDeserializationAfterTimeoutButBeforeExpirationExtension()
+            throws Exception
+    {
+        TestingClock clock = new TestingClock();
+        JweTokenSerializer serializer = tokenSerializer(
+                clock,
+                succinctDuration(12, MINUTES));
+        Date expiration = new Calendar.Builder().setDate(2022, 6, 22).build().getTime();
+        String serializedTokenPair = serializer.serialize(accessAndRefreshTokens("access_token", expiration, "refresh_token"));
+        clock.advanceBy(succinctDuration(10, MINUTES));
+        TokenPair deserializedTokenPair = serializer.deserialize(serializedTokenPair);
+
+        assertThat(deserializedTokenPair.getAccessToken()).isEqualTo("access_token");
+        assertThat(deserializedTokenPair.getExpiration()).isEqualTo(expiration);
+        assertThat(deserializedTokenPair.getRefreshToken()).isEqualTo(Optional.of("refresh_token"));
+    }
+
+    @Test
+    public void testTokenDeserializationAfterTimeoutAndExpirationExtension()
+            throws Exception
+    {
+        TestingClock clock = new TestingClock();
+
+        JweTokenSerializer serializer = tokenSerializer(
+                clock,
+                succinctDuration(12, MINUTES));
+        Date expiration = new Calendar.Builder().setDate(2022, 6, 22).build().getTime();
+        String serializedTokenPair = serializer.serialize(accessAndRefreshTokens("access_token", expiration, "refresh_token"));
+
+        clock.advanceBy(succinctDuration(20, MINUTES));
+        assertThatThrownBy(() -> serializer.deserialize(serializedTokenPair))
+                .isExactlyInstanceOf(ExpiredJwtException.class);
+    }
+
+    private JweTokenSerializer tokenSerializer(Clock clock, Duration tokenExpiration)
+            throws GeneralSecurityException, KeyLengthException
+    {
+        return new JweTokenSerializer(
+                new RefreshTokensConfig(),
+                new Oauth2ClientStub(),
+                "presto_coordinator_test_version",
+                "presto_coordinator",
+                "sub",
+                clock,
+                tokenExpiration);
+    }
+
+    static class Oauth2ClientStub
+            implements OAuth2Client
+    {
+        private final Map claims = Jwts.claims()
+                .setSubject("user");
+
+        @Override
+        public void load()
+        {
+        }
+
+        @Override
+        public Request createAuthorizationRequest(String state, URI callbackUri)
+        {
+            throw new UnsupportedOperationException("operation is not yet supported");
+        }
+
+        @Override
+        public Response getOAuth2Response(String code, URI callbackUri, Optional nonce)
+        {
+            throw new UnsupportedOperationException("operation is not yet supported");
+        }
+
+        @Override
+        public Optional> getClaims(String accessToken)
+        {
+            return Optional.of(claims);
+        }
+
+        @Override
+        public Response refreshTokens(String refreshToken)
+        {
+            throw new UnsupportedOperationException("operation is not yet supported");
+        }
+    }
+
+    private static class TestingClock
+            extends Clock
+    {
+        private Instant currentTime = ZonedDateTime.of(2022, 5, 6, 10, 15, 0, 0, ZoneId.systemDefault()).toInstant();
+
+        @Override
+        public ZoneId getZone()
+        {
+            return ZoneId.systemDefault();
+        }
+
+        @Override
+        public Clock withZone(ZoneId zone)
+        {
+            return this;
+        }
+
+        @Override
+        public Instant instant()
+        {
+            return currentTime;
+        }
+
+        public void advanceBy(Duration currentTimeDelta)
+        {
+            this.currentTime = currentTime.plus(currentTimeDelta.toMillis(), MILLIS);
+        }
+    }
+}
diff --git a/presto-main/src/test/java/com/facebook/presto/server/security/oauth2/TestOAuth2AuthenticationFilterWithJwt.java b/presto-main/src/test/java/com/facebook/presto/server/security/oauth2/TestOAuth2AuthenticationFilterWithJwt.java
new file mode 100644
index 0000000000000..71806b590bc00
--- /dev/null
+++ b/presto-main/src/test/java/com/facebook/presto/server/security/oauth2/TestOAuth2AuthenticationFilterWithJwt.java
@@ -0,0 +1,57 @@
+/*
+ * Licensed 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 com.facebook.presto.server.security.oauth2;
+
+import com.google.common.collect.ImmutableMap;
+import com.google.common.io.Resources;
+
+import java.io.IOException;
+
+import static java.io.File.createTempFile;
+
+public class TestOAuth2AuthenticationFilterWithJwt
+        extends BaseOAuth2AuthenticationFilterTest
+{
+    @Override
+    protected ImmutableMap getOAuth2Config(String idpUrl)
+            throws IOException
+    {
+        return ImmutableMap.builder()
+                .put("http-server.authentication.allow-forwarded-https", "true")
+                .put("http-server.authentication.type", "OAUTH2")
+                .put("http-server.authentication.oauth2.issuer", "https://localhost:4444/")
+                .put("http-server.authentication.oauth2.auth-url", idpUrl + "/oauth2/auth")
+                .put("http-server.authentication.oauth2.token-url", idpUrl + "/oauth2/token")
+                .put("http-server.authentication.oauth2.jwks-url", idpUrl + "/.well-known/jwks.json")
+                .put("http-server.authentication.oauth2.client-id", PRESTO_CLIENT_ID)
+                .put("http-server.authentication.oauth2.client-secret", PRESTO_CLIENT_SECRET)
+                .put("http-server.authentication.oauth2.additional-audiences", TRUSTED_CLIENT_ID)
+                .put("http-server.authentication.oauth2.max-clock-skew", "0s")
+                .put("http-server.authentication.oauth2.user-mapping.pattern", "(.*)(@.*)?")
+                .put("http-server.authentication.oauth2.oidc.discovery", "false")
+                .put("configuration-based-authorizer.role-regex-map.file-path", createTempFile("regex-map", null).getAbsolutePath().toString())
+                .put("oauth2-jwk.http-client.trust-store-path", Resources.getResource("cert/localhost.pem").getPath())
+                .build();
+    }
+
+    @Override
+    protected TestingHydraIdentityProvider getHydraIdp()
+            throws Exception
+    {
+        TestingHydraIdentityProvider hydraIdP = new TestingHydraIdentityProvider(TTL_ACCESS_TOKEN_IN_SECONDS, true, false);
+        hydraIdP.start();
+
+        return hydraIdP;
+    }
+}
diff --git a/presto-main/src/test/java/com/facebook/presto/server/security/oauth2/TestOAuth2AuthenticationFilterWithOpaque.java b/presto-main/src/test/java/com/facebook/presto/server/security/oauth2/TestOAuth2AuthenticationFilterWithOpaque.java
new file mode 100644
index 0000000000000..114b9b0ed20d1
--- /dev/null
+++ b/presto-main/src/test/java/com/facebook/presto/server/security/oauth2/TestOAuth2AuthenticationFilterWithOpaque.java
@@ -0,0 +1,60 @@
+/*
+ * Licensed 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 com.facebook.presto.server.security.oauth2;
+
+import com.google.common.collect.ImmutableMap;
+import com.google.common.io.Resources;
+
+import java.io.IOException;
+
+import static java.io.File.createTempFile;
+
+public class TestOAuth2AuthenticationFilterWithOpaque
+        extends BaseOAuth2AuthenticationFilterTest
+{
+    @Override
+    protected ImmutableMap getOAuth2Config(String idpUrl)
+            throws IOException
+    {
+        return ImmutableMap.builder()
+                .put("http-server.authentication.allow-forwarded-https", "true")
+                .put("http-server.authentication.type", "OAUTH2")
+                .put("http-server.authentication.oauth2.issuer", "https://localhost:4444/")
+                .put("http-server.authentication.oauth2.auth-url", idpUrl + "/oauth2/auth")
+                .put("http-server.authentication.oauth2.token-url", idpUrl + "/oauth2/token")
+                .put("http-server.authentication.oauth2.jwks-url", idpUrl + "/.well-known/jwks.json")
+                .put("http-server.authentication.oauth2.userinfo-url", idpUrl + "/userinfo")
+                .put("http-server.authentication.oauth2.client-id", PRESTO_CLIENT_ID)
+                .put("http-server.authentication.oauth2.client-secret", PRESTO_CLIENT_SECRET)
+                // This is necessary as Hydra does not return `sub` from `/userinfo` for client credential grants.
+                .put("http-server.authentication.oauth2.principal-field", "iss")
+                .put("http-server.authentication.oauth2.additional-audiences", TRUSTED_CLIENT_ID)
+                .put("http-server.authentication.oauth2.max-clock-skew", "0s")
+                .put("http-server.authentication.oauth2.user-mapping.pattern", "(.*)(@.*)?")
+                .put("http-server.authentication.oauth2.oidc.discovery", "false")
+                .put("configuration-based-authorizer.role-regex-map.file-path", createTempFile("regex-map", null).getAbsolutePath().toString())
+                .put("oauth2-jwk.http-client.trust-store-path", Resources.getResource("cert/localhost.pem").getPath())
+                .build();
+    }
+
+    @Override
+    protected TestingHydraIdentityProvider getHydraIdp()
+            throws Exception
+    {
+        TestingHydraIdentityProvider hydraIdP = new TestingHydraIdentityProvider(TTL_ACCESS_TOKEN_IN_SECONDS, false, false);
+        hydraIdP.start();
+
+        return hydraIdP;
+    }
+}
diff --git a/presto-main/src/test/java/com/facebook/presto/server/security/oauth2/TestOAuth2Config.java b/presto-main/src/test/java/com/facebook/presto/server/security/oauth2/TestOAuth2Config.java
new file mode 100644
index 0000000000000..bb2735b95e7f6
--- /dev/null
+++ b/presto-main/src/test/java/com/facebook/presto/server/security/oauth2/TestOAuth2Config.java
@@ -0,0 +1,93 @@
+/*
+ * Licensed 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 com.facebook.presto.server.security.oauth2;
+
+import com.facebook.airlift.units.Duration;
+import com.google.common.collect.ImmutableMap;
+import org.testng.annotations.Test;
+
+import java.io.IOException;
+import java.nio.file.Files;
+import java.nio.file.Path;
+import java.util.Map;
+
+import static com.facebook.airlift.configuration.testing.ConfigAssertions.assertFullMapping;
+import static com.facebook.airlift.configuration.testing.ConfigAssertions.assertRecordedDefaults;
+import static com.facebook.airlift.configuration.testing.ConfigAssertions.recordDefaults;
+import static java.util.concurrent.TimeUnit.MINUTES;
+import static java.util.concurrent.TimeUnit.SECONDS;
+
+public class TestOAuth2Config
+{
+    @Test
+    public void testDefaults()
+    {
+        assertRecordedDefaults(recordDefaults(OAuth2Config.class)
+                .setStateKey(null)
+                .setIssuer(null)
+                .setClientId(null)
+                .setClientSecret(null)
+                .setScopes("openid")
+                .setChallengeTimeout(new Duration(15, MINUTES))
+                .setPrincipalField("sub")
+                .setGroupsField(null)
+                .setAdditionalAudiences("")
+                .setMaxClockSkew(new Duration(1, MINUTES))
+                .setUserMappingPattern(null)
+                .setUserMappingFile(null)
+                .setEnableRefreshTokens(false)
+                .setEnableDiscovery(true));
+    }
+
+    @Test
+    public void testExplicitPropertyMappings()
+            throws IOException
+    {
+        Path userMappingFile = Files.createTempFile(null, null);
+        Map properties = ImmutableMap.builder()
+                .put("http-server.authentication.oauth2.state-key", "key-secret")
+                .put("http-server.authentication.oauth2.issuer", "http://127.0.0.1:9000/oauth2")
+                .put("http-server.authentication.oauth2.client-id", "another-consumer")
+                .put("http-server.authentication.oauth2.client-secret", "consumer-secret")
+                .put("http-server.authentication.oauth2.scopes", "email,offline")
+                .put("http-server.authentication.oauth2.principal-field", "some-field")
+                .put("http-server.authentication.oauth2.groups-field", "groups")
+                .put("http-server.authentication.oauth2.additional-audiences", "test-aud1,test-aud2")
+                .put("http-server.authentication.oauth2.challenge-timeout", "90s")
+                .put("http-server.authentication.oauth2.max-clock-skew", "15s")
+                .put("http-server.authentication.oauth2.user-mapping.pattern", "(.*)@something")
+                .put("http-server.authentication.oauth2.user-mapping.file", userMappingFile.toString())
+                .put("http-server.authentication.oauth2.refresh-tokens", "true")
+                .put("http-server.authentication.oauth2.oidc.discovery", "false")
+                .build();
+
+        OAuth2Config expected = new OAuth2Config()
+                .setStateKey("key-secret")
+                .setIssuer("http://127.0.0.1:9000/oauth2")
+                .setClientId("another-consumer")
+                .setClientSecret("consumer-secret")
+                .setScopes("email, offline")
+                .setPrincipalField("some-field")
+                .setGroupsField("groups")
+                .setAdditionalAudiences("test-aud1,test-aud2")
+                .setChallengeTimeout(new Duration(90, SECONDS))
+                .setMaxClockSkew(new Duration(15, SECONDS))
+                .setUserMappingPattern("(.*)@something")
+                .setUserMappingFile(userMappingFile.toFile())
+                .setEnableRefreshTokens(true)
+                .setEnableDiscovery(false);
+
+        assertFullMapping(properties, expected);
+    }
+}
diff --git a/presto-main/src/test/java/com/facebook/presto/server/security/oauth2/TestOAuth2Utils.java b/presto-main/src/test/java/com/facebook/presto/server/security/oauth2/TestOAuth2Utils.java
new file mode 100644
index 0000000000000..119fd7213bd96
--- /dev/null
+++ b/presto-main/src/test/java/com/facebook/presto/server/security/oauth2/TestOAuth2Utils.java
@@ -0,0 +1,54 @@
+/*
+ * Licensed 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 com.facebook.presto.server.security.oauth2;
+
+import com.facebook.presto.server.MockHttpServletRequest;
+import com.google.common.collect.ImmutableListMultimap;
+import jakarta.servlet.http.HttpServletRequest;
+import jakarta.ws.rs.core.UriBuilder;
+import org.testng.annotations.Test;
+
+import static com.facebook.presto.server.security.oauth2.OAuth2Utils.getSchemeUriBuilder;
+import static com.google.common.net.HttpHeaders.X_FORWARDED_PROTO;
+import static org.testng.Assert.assertEquals;
+
+public class TestOAuth2Utils
+{
+    @Test
+    public void testGetSchemeUriBuilderNoProtoHeader()
+    {
+        HttpServletRequest request = new MockHttpServletRequest(
+                ImmutableListMultimap.builder()
+                        .build(),
+                "testRemote",
+                   "http://www.example.com");
+
+        UriBuilder builder = getSchemeUriBuilder(request);
+        assertEquals(builder.build().getScheme(), "http");
+    }
+
+    @Test
+    public void testGetSchemeUriBuilderProtoHeader()
+    {
+        HttpServletRequest request = new MockHttpServletRequest(
+                ImmutableListMultimap.builder()
+                        .put(X_FORWARDED_PROTO, "https")
+                        .build(),
+                "testRemote",
+                "http://www.example.com");
+
+        UriBuilder builder = getSchemeUriBuilder(request);
+        assertEquals(builder.build().getScheme(), "https");
+    }
+}
diff --git a/presto-main/src/test/java/com/facebook/presto/server/security/oauth2/TestOidcDiscovery.java b/presto-main/src/test/java/com/facebook/presto/server/security/oauth2/TestOidcDiscovery.java
new file mode 100644
index 0000000000000..679095b9c7cd4
--- /dev/null
+++ b/presto-main/src/test/java/com/facebook/presto/server/security/oauth2/TestOidcDiscovery.java
@@ -0,0 +1,373 @@
+/*
+ * Licensed 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 com.facebook.presto.server.security.oauth2;
+
+import com.facebook.airlift.http.server.Authenticator;
+import com.facebook.airlift.http.server.HttpServerConfig;
+import com.facebook.airlift.http.server.HttpServerInfo;
+import com.facebook.airlift.http.server.testing.TestingHttpServer;
+import com.facebook.airlift.node.NodeInfo;
+import com.facebook.presto.server.security.oauth2.OAuth2ServerConfigProvider.OAuth2ServerConfig;
+import com.facebook.presto.server.testing.TestingPrestoServer;
+import com.google.common.collect.ImmutableMap;
+import com.google.common.io.Resources;
+import com.google.inject.Key;
+import com.google.inject.TypeLiteral;
+import jakarta.servlet.http.HttpServlet;
+import jakarta.servlet.http.HttpServletRequest;
+import jakarta.servlet.http.HttpServletResponse;
+import org.testng.annotations.DataProvider;
+import org.testng.annotations.Test;
+
+import java.io.IOException;
+import java.net.URI;
+import java.time.Instant;
+import java.util.List;
+import java.util.Map;
+import java.util.Optional;
+
+import static com.facebook.airlift.http.client.HttpStatus.TOO_MANY_REQUESTS;
+import static com.facebook.presto.server.security.oauth2.BaseOAuth2AuthenticationFilterTest.PRESTO_CLIENT_ID;
+import static com.facebook.presto.server.security.oauth2.BaseOAuth2AuthenticationFilterTest.PRESTO_CLIENT_SECRET;
+import static jakarta.ws.rs.core.HttpHeaders.CONTENT_TYPE;
+import static jakarta.ws.rs.core.MediaType.APPLICATION_JSON;
+import static java.io.File.createTempFile;
+import static java.nio.charset.StandardCharsets.UTF_8;
+import static java.util.Objects.requireNonNull;
+import static org.assertj.core.api.Assertions.assertThat;
+import static org.assertj.core.api.Assertions.assertThatThrownBy;
+
+public class TestOidcDiscovery
+{
+    @Test(dataProvider = "staticConfiguration")
+    public void testStaticConfiguration(Optional accessTokenPath, Optional userinfoPath)
+            throws Exception
+    {
+        try (MetadataServer metadataServer = new MetadataServer(ImmutableMap.of("/jwks.json", "jwk/jwk-public.json"))) {
+            URI issuer = metadataServer.getBaseUrl();
+            Optional accessTokenIssuer = accessTokenPath.map(issuer::resolve);
+            Optional userinfoUrl = userinfoPath.map(issuer::resolve);
+            ImmutableMap.Builder properties = ImmutableMap.builder()
+                    .put("http-server.authentication.oauth2.issuer", metadataServer.getBaseUrl().toString())
+                    .put("http-server.authentication.oauth2.oidc.discovery", "false")
+                    .put("http-server.authentication.oauth2.auth-url", issuer.resolve("/connect/authorize").toString())
+                    .put("http-server.authentication.oauth2.token-url", issuer.resolve("/connect/token").toString())
+                    .put("http-server.authentication.oauth2.jwks-url", issuer.resolve("/jwks.json").toString())
+                    .put("configuration-based-authorizer.role-regex-map.file-path", createTempFile("regex-map", null).getAbsolutePath().toString());
+            accessTokenIssuer.map(URI::toString).ifPresent(uri -> properties.put("http-server.authentication.oauth2.access-token-issuer", uri));
+            userinfoUrl.map(URI::toString).ifPresent(uri -> properties.put("http-server.authentication.oauth2.userinfo-url", uri));
+            try (TestingPrestoServer server = createServer(properties.build())) {
+                assertConfiguration(server, issuer, accessTokenIssuer.map(issuer::resolve), userinfoUrl.map(issuer::resolve));
+            }
+        }
+    }
+
+    @DataProvider(name = "staticConfiguration")
+    public static Object[][] staticConfiguration()
+    {
+        return new Object[][] {
+                {Optional.empty(), Optional.empty()},
+                {Optional.of("/access-token-issuer"), Optional.of("/userinfo")},
+        };
+    }
+
+    @Test(dataProvider = "oidcDiscovery")
+    public void testOidcDiscovery(String configuration, Optional accessTokenIssuer, Optional userinfoUrl)
+            throws Exception
+    {
+        try (MetadataServer metadataServer = new MetadataServer(
+                ImmutableMap.builder()
+                        .put("/.well-known/openid-configuration", "oidc/" + configuration)
+                        .put("/jwks.json", "jwk/jwk-public.json")
+                        .build());
+                TestingPrestoServer server = createServer(
+                        ImmutableMap.builder()
+                                .put("http-server.authentication.oauth2.issuer", metadataServer.getBaseUrl().toString())
+                                .put("http-server.authentication.oauth2.oidc.discovery", "true")
+                                .put("configuration-based-authorizer.role-regex-map.file-path", createTempFile("regex-map", null).getAbsolutePath().toString())
+                                .build())) {
+            URI issuer = metadataServer.getBaseUrl();
+            assertConfiguration(server, issuer, accessTokenIssuer.map(issuer::resolve), userinfoUrl.map(issuer::resolve));
+        }
+    }
+
+    @DataProvider(name = "oidcDiscovery")
+    public static Object[][] oidcDiscovery()
+    {
+        return new Object[][] {
+                {"openid-configuration.json", Optional.empty(), Optional.of("/connect/userinfo")},
+                {"openid-configuration-without-userinfo.json", Optional.empty(), Optional.empty()},
+                {"openid-configuration-with-access-token-issuer.json", Optional.of("http://access-token-issuer.com/adfs/services/trust"), Optional.of("/connect/userinfo")},
+        };
+    }
+
+    @Test
+    public void testIssuerCheck()
+    {
+        assertThatThrownBy(() -> {
+            try (MetadataServer metadataServer = new MetadataServer(
+                    ImmutableMap.builder()
+                            .put("/.well-known/openid-configuration", "oidc/openid-configuration-invalid-issuer.json")
+                            .put("/jwks.json", "jwk/jwk-public.json")
+                            .build());
+                    TestingPrestoServer server = createServer(
+                            ImmutableMap.builder()
+                                    .put("http-server.authentication.oauth2.issuer", metadataServer.getBaseUrl().toString())
+                                    .put("http-server.authentication.oauth2.oidc.discovery", "true")
+                                    .put("configuration-based-authorizer.role-regex-map.file-path", createTempFile("regex-map", null).getAbsolutePath().toString())
+                                    .build())) {
+                // should throw an exception
+                server.getInstance(Key.get(OAuth2ServerConfigProvider.class)).get();
+            }
+        }).hasMessageContaining(
+                "Invalid response from OpenID Metadata endpoint. " +
+                        "The value of the \"issuer\" claim in Metadata document different than the Issuer URL used for the Configuration Request.");
+    }
+
+    @Test
+    public void testStopOnClientError()
+    {
+        assertThatThrownBy(() -> {
+            try (MetadataServer metadataServer = new MetadataServer(ImmutableMap.of());
+                    TestingPrestoServer server = createServer(
+                            ImmutableMap.builder()
+                                    .put("http-server.authentication.oauth2.issuer", metadataServer.getBaseUrl().toString())
+                                    .put("http-server.authentication.oauth2.oidc.discovery", "true")
+                                    .put("configuration-based-authorizer.role-regex-map.file-path", createTempFile("regex-map", null).getAbsolutePath().toString())
+                                    .build())) {
+                // should throw an exception
+                server.getInstance(Key.get(OAuth2ServerConfigProvider.class)).get();
+            }
+        }).hasMessageContaining("Invalid response from OpenID Metadata endpoint. Expected response code to be 200, but was 404");
+    }
+
+    @Test
+    public void testOidcDiscoveryRetrying()
+            throws Exception
+    {
+        try (MetadataServer metadataServer = new MetadataServer(new MetadataServletWithStartup(
+                ImmutableMap.builder()
+                        .put("/.well-known/openid-configuration", "oidc/openid-configuration.json")
+                        .put("/jwks.json", "jwk/jwk-public.json")
+                        .build(), 5));
+                TestingPrestoServer server = createServer(
+                        ImmutableMap.builder()
+                                .put("http-server.authentication.oauth2.issuer", metadataServer.getBaseUrl().toString())
+                                .put("http-server.authentication.oauth2.oidc.discovery", "true")
+                                .put("http-server.authentication.oauth2.oidc.discovery.timeout", "10s")
+                                .put("configuration-based-authorizer.role-regex-map.file-path", createTempFile("regex-map", null).getAbsolutePath().toString())
+                                .build())) {
+            URI issuer = metadataServer.getBaseUrl();
+            assertConfiguration(server, issuer, Optional.empty(), Optional.of(issuer.resolve("/connect/userinfo")));
+        }
+    }
+
+    @Test
+    public void testOidcDiscoveryTimesOut()
+    {
+        assertThatThrownBy(() -> {
+            try (MetadataServer metadataServer = new MetadataServer(new MetadataServletWithStartup(
+                    ImmutableMap.builder()
+                            .put("/.well-known/openid-configuration", "oidc/openid-configuration.json")
+                            .put("/jwks.json", "jwk/jwk-public.json")
+                            .build(), 10));
+                    TestingPrestoServer server = createServer(
+                            ImmutableMap.builder()
+                                    .put("http-server.authentication.oauth2.issuer", metadataServer.getBaseUrl().toString())
+                                    .put("http-server.authentication.oauth2.oidc.discovery", "true")
+                                    .put("http-server.authentication.oauth2.oidc.discovery.timeout", "5s")
+                                    .put("configuration-based-authorizer.role-regex-map.file-path", createTempFile("regex-map", null).getAbsolutePath().toString())
+                                    .build())) {
+                // should throw an exception
+                server.getInstance(Key.get(OAuth2ServerConfigProvider.class)).get();
+            }
+        }).hasMessageContaining("Invalid response from OpenID Metadata endpoint: 429");
+    }
+
+    @Test
+    public void testIgnoringUserinfoUrl()
+            throws Exception
+    {
+        try (MetadataServer metadataServer = new MetadataServer(
+                ImmutableMap.builder()
+                        .put("/.well-known/openid-configuration", "oidc/openid-configuration.json")
+                        .put("/jwks.json", "jwk/jwk-public.json")
+                        .build());
+                TestingPrestoServer server = createServer(
+                        ImmutableMap.builder()
+                                .put("http-server.authentication.oauth2.issuer", metadataServer.getBaseUrl().toString())
+                                .put("http-server.authentication.oauth2.oidc.discovery", "true")
+                                .put("http-server.authentication.oauth2.oidc.use-userinfo-endpoint", "false")
+                                .put("configuration-based-authorizer.role-regex-map.file-path", createTempFile("regex-map", null).getAbsolutePath().toString())
+                                .build())) {
+            URI issuer = metadataServer.getBaseUrl();
+            assertConfiguration(server, issuer, Optional.empty(), Optional.empty());
+        }
+    }
+
+    @Test
+    public void testBackwardCompatibility()
+            throws Exception
+    {
+        try (MetadataServer metadataServer = new MetadataServer(
+                ImmutableMap.builder()
+                        .put("/.well-known/openid-configuration", "oidc/openid-configuration-with-access-token-issuer.json")
+                        .put("/jwks.json", "jwk/jwk-public.json")
+                        .build())) {
+            URI issuer = metadataServer.getBaseUrl();
+            URI authUrl = issuer.resolve("/custom-authorize");
+            URI tokenUrl = issuer.resolve("/custom-token");
+            URI jwksUrl = issuer.resolve("/custom-jwks.json");
+            String accessTokenIssuer = issuer.resolve("/custom-access-token-issuer").toString();
+            URI userinfoUrl = issuer.resolve("/custom-userinfo-url");
+            try (TestingPrestoServer server = createServer(
+                    ImmutableMap.builder()
+                            .put("http-server.authentication.oauth2.issuer", issuer.toString())
+                            .put("http-server.authentication.oauth2.oidc.discovery", "true")
+                            .put("http-server.authentication.oauth2.auth-url", authUrl.toString())
+                            .put("http-server.authentication.oauth2.token-url", tokenUrl.toString())
+                            .put("http-server.authentication.oauth2.jwks-url", jwksUrl.toString())
+                            .put("http-server.authentication.oauth2.access-token-issuer", accessTokenIssuer)
+                            .put("http-server.authentication.oauth2.userinfo-url", userinfoUrl.toString())
+                            .put("configuration-based-authorizer.role-regex-map.file-path", createTempFile("regex-map", null).getAbsolutePath().toString())
+                            .build())) {
+                assertComponents(server);
+                OAuth2ServerConfig config = server.getInstance(Key.get(OAuth2ServerConfigProvider.class)).get();
+                assertThat(config.getAccessTokenIssuer()).isEqualTo(Optional.of(accessTokenIssuer));
+                assertThat(config.getAuthUrl()).isEqualTo(authUrl);
+                assertThat(config.getTokenUrl()).isEqualTo(tokenUrl);
+                assertThat(config.getJwksUrl()).isEqualTo(jwksUrl);
+                assertThat(config.getUserinfoUrl()).isEqualTo(Optional.of(userinfoUrl));
+            }
+        }
+    }
+
+    private static void assertConfiguration(TestingPrestoServer server, URI issuer, Optional accessTokenIssuer, Optional userinfoUrl)
+    {
+        assertComponents(server);
+        OAuth2ServerConfig config = server.getInstance(Key.get(OAuth2ServerConfigProvider.class)).get();
+        assertThat(config.getAccessTokenIssuer()).isEqualTo(accessTokenIssuer.map(URI::toString));
+        assertThat(config.getAuthUrl()).isEqualTo(issuer.resolve("/connect/authorize"));
+        assertThat(config.getTokenUrl()).isEqualTo(issuer.resolve("/connect/token"));
+        assertThat(config.getJwksUrl()).isEqualTo(issuer.resolve("/jwks.json"));
+        assertThat(config.getUserinfoUrl()).isEqualTo(userinfoUrl);
+    }
+
+    private static void assertComponents(TestingPrestoServer server)
+    {
+        List authenticators = server.getInstance(Key.get(new TypeLiteral>() {}));
+        assertThat(authenticators).hasSize(1);
+        assertThat(authenticators.get(0)).isInstanceOf(OAuth2Authenticator.class);
+//        assertThat(server.getInstance(Key.get(WebUiAuthenticationFilter.class))).isInstanceOf(OAuth2WebUiAuthenticationFilter.class);
+        // does not throw an exception
+        server.getInstance(Key.get(OAuth2Client.class)).load();
+    }
+
+    private static TestingPrestoServer createServer(Map configuration)
+            throws Exception
+    {
+        ImmutableMap config = ImmutableMap.builder()
+                .put("http-server.authentication.allow-forwarded-https", "true")
+                .put("http-server.authentication.type", "OAUTH2")
+                .put("http-server.authentication.oauth2.client-id", PRESTO_CLIENT_ID)
+                .put("http-server.authentication.oauth2.client-secret", PRESTO_CLIENT_SECRET)
+                .putAll(configuration)
+                .build();
+
+        return new TestingPrestoServer(config);
+    }
+
+    public static class MetadataServer
+            implements AutoCloseable
+    {
+        private final TestingHttpServer httpServer;
+
+        public MetadataServer(Map responseMapping)
+                throws Exception
+        {
+            this(new MetadataServlet(responseMapping));
+        }
+
+        public MetadataServer(HttpServlet servlet)
+                throws Exception
+        {
+            NodeInfo nodeInfo = new NodeInfo("test");
+            HttpServerConfig config = new HttpServerConfig().setHttpPort(0);
+            HttpServerInfo httpServerInfo = new HttpServerInfo(config, nodeInfo);
+            httpServer = new TestingHttpServer(httpServerInfo, nodeInfo, config, servlet, ImmutableMap.of(), ImmutableMap.of(), Optional.empty());
+            httpServer.start();
+        }
+
+        public URI getBaseUrl()
+        {
+            return httpServer.getBaseUrl();
+        }
+
+        @Override
+        public void close()
+                throws Exception
+        {
+            httpServer.stop();
+        }
+    }
+
+    public static class MetadataServlet
+            extends HttpServlet
+    {
+        private final Map responseMapping;
+
+        public MetadataServlet(Map responseMapping)
+        {
+            this.responseMapping = requireNonNull(responseMapping, "responseMapping is null");
+        }
+
+        @Override
+        protected void doGet(HttpServletRequest request, HttpServletResponse response)
+                throws IOException
+        {
+            String fileName = responseMapping.get(request.getPathInfo());
+            if (fileName == null) {
+                response.setStatus(404);
+                return;
+            }
+            response.setHeader(CONTENT_TYPE, APPLICATION_JSON);
+            String body = Resources.toString(Resources.getResource(fileName), UTF_8);
+            body = body.replaceAll("https://issuer.com", request.getRequestURL().toString().replace("/.well-known/openid-configuration", ""));
+            response.getWriter().write(body);
+        }
+    }
+
+    public static class MetadataServletWithStartup
+            extends MetadataServlet
+    {
+        private final Instant startTime;
+
+        public MetadataServletWithStartup(Map responseMapping, int startupInSeconds)
+        {
+            super(responseMapping);
+            startTime = Instant.now().plusSeconds(startupInSeconds);
+        }
+
+        @Override
+        protected void doGet(HttpServletRequest request, HttpServletResponse response)
+                throws IOException
+        {
+            if (Instant.now().isBefore(startTime)) {
+                response.setStatus(TOO_MANY_REQUESTS.code());
+                return;
+            }
+            super.doGet(request, response);
+        }
+    }
+}
diff --git a/presto-main/src/test/java/com/facebook/presto/server/security/oauth2/TestOidcDiscoveryConfig.java b/presto-main/src/test/java/com/facebook/presto/server/security/oauth2/TestOidcDiscoveryConfig.java
new file mode 100644
index 0000000000000..37c4dd317924b
--- /dev/null
+++ b/presto-main/src/test/java/com/facebook/presto/server/security/oauth2/TestOidcDiscoveryConfig.java
@@ -0,0 +1,67 @@
+/*
+ * Licensed 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 com.facebook.presto.server.security.oauth2;
+
+import com.facebook.airlift.units.Duration;
+import com.google.common.collect.ImmutableMap;
+import org.testng.annotations.Test;
+
+import java.util.Map;
+
+import static com.facebook.airlift.configuration.testing.ConfigAssertions.assertFullMapping;
+import static com.facebook.airlift.configuration.testing.ConfigAssertions.assertRecordedDefaults;
+import static com.facebook.airlift.configuration.testing.ConfigAssertions.recordDefaults;
+import static java.util.concurrent.TimeUnit.MINUTES;
+import static java.util.concurrent.TimeUnit.SECONDS;
+
+public class TestOidcDiscoveryConfig
+{
+    @Test
+    public void testDefaults()
+    {
+        assertRecordedDefaults(recordDefaults(OidcDiscoveryConfig.class)
+                .setDiscoveryTimeout(new Duration(30, SECONDS))
+                .setUserinfoEndpointEnabled(true)
+                .setAccessTokenIssuer(null)
+                .setAuthUrl(null)
+                .setTokenUrl(null)
+                .setJwksUrl(null)
+                .setUserinfoUrl(null));
+    }
+
+    @Test
+    public void testExplicitPropertyMapping()
+    {
+        Map properties = ImmutableMap.builder()
+                .put("http-server.authentication.oauth2.oidc.discovery.timeout", "1m")
+                .put("http-server.authentication.oauth2.oidc.use-userinfo-endpoint", "false")
+                .put("http-server.authentication.oauth2.access-token-issuer", "https://issuer.com/at")
+                .put("http-server.authentication.oauth2.auth-url", "https://issuer.com/auth")
+                .put("http-server.authentication.oauth2.token-url", "https://issuer.com/token")
+                .put("http-server.authentication.oauth2.jwks-url", "https://issuer.com/jwks.json")
+                .put("http-server.authentication.oauth2.userinfo-url", "https://issuer.com/user")
+                .build();
+
+        OidcDiscoveryConfig expected = new OidcDiscoveryConfig()
+                .setDiscoveryTimeout(new Duration(1, MINUTES))
+                .setUserinfoEndpointEnabled(false)
+                .setAccessTokenIssuer("https://issuer.com/at")
+                .setAuthUrl("https://issuer.com/auth")
+                .setTokenUrl("https://issuer.com/token")
+                .setJwksUrl("https://issuer.com/jwks.json")
+                .setUserinfoUrl("https://issuer.com/user");
+
+        assertFullMapping(properties, expected);
+    }
+}
diff --git a/presto-main/src/test/java/com/facebook/presto/server/security/oauth2/TestRefreshTokensConfig.java b/presto-main/src/test/java/com/facebook/presto/server/security/oauth2/TestRefreshTokensConfig.java
new file mode 100644
index 0000000000000..29c7d5b11c603
--- /dev/null
+++ b/presto-main/src/test/java/com/facebook/presto/server/security/oauth2/TestRefreshTokensConfig.java
@@ -0,0 +1,74 @@
+/*
+ * Licensed 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 com.facebook.presto.server.security.oauth2;
+
+import com.google.common.collect.ImmutableMap;
+import org.testng.annotations.Test;
+
+import javax.crypto.KeyGenerator;
+import javax.crypto.SecretKey;
+
+import java.security.NoSuchAlgorithmException;
+import java.util.Map;
+
+import static com.facebook.airlift.configuration.testing.ConfigAssertions.assertFullMapping;
+import static com.facebook.airlift.configuration.testing.ConfigAssertions.assertRecordedDefaults;
+import static com.facebook.airlift.configuration.testing.ConfigAssertions.recordDefaults;
+import static com.facebook.airlift.units.Duration.succinctDuration;
+import static io.jsonwebtoken.io.Encoders.BASE64;
+import static java.util.concurrent.TimeUnit.HOURS;
+
+public class TestRefreshTokensConfig
+{
+    @Test
+    public void testDefaults()
+    {
+        assertRecordedDefaults(recordDefaults(RefreshTokensConfig.class)
+                .setTokenExpiration(succinctDuration(1, HOURS))
+                .setIssuer("Presto_coordinator")
+                .setAudience("Presto_coordinator")
+                .setSecretKey(null));
+    }
+
+    @Test
+    public void testExplicitPropertyMappings()
+            throws Exception
+    {
+        String encodedBase64SecretKey = BASE64.encode(generateKey());
+
+        Map properties = ImmutableMap.builder()
+                .put("http-server.authentication.oauth2.refresh-tokens.issued-token.timeout", "24h")
+                .put("http-server.authentication.oauth2.refresh-tokens.issued-token.issuer", "issuer")
+                .put("http-server.authentication.oauth2.refresh-tokens.issued-token.audience", "audience")
+                .put("http-server.authentication.oauth2.refresh-tokens.secret-key", encodedBase64SecretKey)
+                .build();
+
+        RefreshTokensConfig expected = new RefreshTokensConfig()
+                .setTokenExpiration(succinctDuration(24, HOURS))
+                .setIssuer("issuer")
+                .setAudience("audience")
+                .setSecretKey(encodedBase64SecretKey);
+
+        assertFullMapping(properties, expected);
+    }
+
+    private byte[] generateKey()
+            throws NoSuchAlgorithmException
+    {
+        KeyGenerator generator = KeyGenerator.getInstance("AES");
+        generator.init(256);
+        SecretKey key = generator.generateKey();
+        return key.getEncoded();
+    }
+}
diff --git a/presto-main/src/test/java/com/facebook/presto/server/security/oauth2/TestingHydraIdentityProvider.java b/presto-main/src/test/java/com/facebook/presto/server/security/oauth2/TestingHydraIdentityProvider.java
new file mode 100644
index 0000000000000..0daa570a29842
--- /dev/null
+++ b/presto-main/src/test/java/com/facebook/presto/server/security/oauth2/TestingHydraIdentityProvider.java
@@ -0,0 +1,391 @@
+/*
+ * Licensed 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 com.facebook.presto.server.security.oauth2;
+
+import com.facebook.airlift.http.server.HttpServerConfig;
+import com.facebook.airlift.http.server.HttpServerInfo;
+import com.facebook.airlift.http.server.testing.TestingHttpServer;
+import com.facebook.airlift.log.Level;
+import com.facebook.airlift.log.Logging;
+import com.facebook.airlift.node.NodeInfo;
+import com.facebook.presto.server.testing.TestingPrestoServer;
+import com.facebook.presto.util.AutoCloseableCloser;
+import com.fasterxml.jackson.databind.JsonNode;
+import com.fasterxml.jackson.databind.ObjectMapper;
+import com.fasterxml.jackson.databind.node.ObjectNode;
+import com.google.common.collect.ImmutableList;
+import com.google.common.collect.ImmutableMap;
+import com.google.common.io.Resources;
+import com.google.inject.Key;
+import com.nimbusds.oauth2.sdk.GrantType;
+import jakarta.servlet.http.HttpServlet;
+import jakarta.servlet.http.HttpServletRequest;
+import jakarta.servlet.http.HttpServletResponse;
+import jakarta.ws.rs.core.HttpHeaders;
+import okhttp3.Credentials;
+import okhttp3.FormBody;
+import okhttp3.HttpUrl;
+import okhttp3.MediaType;
+import okhttp3.OkHttpClient;
+import okhttp3.Request;
+import okhttp3.RequestBody;
+import okhttp3.Response;
+import org.testcontainers.containers.FixedHostPortGenericContainer;
+import org.testcontainers.containers.GenericContainer;
+import org.testcontainers.containers.Network;
+import org.testcontainers.containers.PostgreSQLContainer;
+import org.testcontainers.containers.startupcheck.OneShotStartupCheckStrategy;
+import org.testcontainers.containers.wait.strategy.Wait;
+import org.testcontainers.containers.wait.strategy.WaitAllStrategy;
+import org.testcontainers.utility.MountableFile;
+
+import java.io.Closeable;
+import java.io.IOException;
+import java.net.URI;
+import java.time.Duration;
+import java.util.List;
+import java.util.Optional;
+
+import static com.facebook.presto.client.OkHttpUtil.setupInsecureSsl;
+import static com.facebook.presto.server.security.oauth2.TokenEndpointAuthMethod.CLIENT_SECRET_BASIC;
+import static com.google.common.base.Preconditions.checkState;
+import static com.google.common.base.Throwables.throwIfUnchecked;
+import static jakarta.servlet.http.HttpServletResponse.SC_NOT_FOUND;
+import static jakarta.servlet.http.HttpServletResponse.SC_OK;
+import static jakarta.ws.rs.core.MediaType.APPLICATION_JSON;
+import static java.util.Objects.requireNonNull;
+
+public class TestingHydraIdentityProvider
+        implements Closeable
+{
+    private static final String HYDRA_IMAGE = "oryd/hydra:v1.10.6";
+    private static final String ISSUER = "https://localhost:4444/";
+    private static final String DSN = "postgres://hydra:mysecretpassword@database:5432/hydra?sslmode=disable";
+
+    private final Network network = Network.newNetwork();
+
+    private final PostgreSQLContainer databaseContainer = new PostgreSQLContainer<>()
+            .withNetwork(network)
+            .withNetworkAliases("database")
+            .withUsername("hydra")
+            .withPassword("mysecretpassword")
+            .withDatabaseName("hydra");
+
+    private final GenericContainer migrationContainer = createHydraContainer()
+            .withCommand("migrate", "sql", "--yes", DSN)
+            .withStartupCheckStrategy(new OneShotStartupCheckStrategy().withTimeout(Duration.ofMinutes(5)));
+
+    private final AutoCloseableCloser closer = AutoCloseableCloser.create();
+    private final ObjectMapper mapper = new ObjectMapper();
+    private final Duration ttlAccessToken;
+    private final boolean useJwt;
+    private final boolean exposeFixedPorts;
+    private final OkHttpClient httpClient;
+    private FixedHostPortGenericContainer hydraContainer;
+
+    public TestingHydraIdentityProvider()
+    {
+        this(Duration.ofMinutes(30), true, false);
+    }
+
+    public TestingHydraIdentityProvider(Duration ttlAccessToken, boolean useJwt, boolean exposeFixedPorts)
+    {
+        this.ttlAccessToken = requireNonNull(ttlAccessToken, "ttlAccessToken is null");
+        this.useJwt = useJwt;
+        this.exposeFixedPorts = exposeFixedPorts;
+        OkHttpClient.Builder httpClientBuilder = new OkHttpClient.Builder();
+        setupInsecureSsl(httpClientBuilder);
+        httpClientBuilder.followRedirects(false);
+        httpClient = httpClientBuilder.build();
+        closer.register(network);
+        closer.register(databaseContainer);
+        closer.register(migrationContainer);
+    }
+
+    public void start()
+            throws Exception
+    {
+        databaseContainer.start();
+        migrationContainer.start();
+        TestingHttpServer loginAndConsentServer = createTestingLoginAndConsentServer();
+        closer.register(loginAndConsentServer::stop);
+        loginAndConsentServer.start();
+        URI loginAndConsentBaseUrl = loginAndConsentServer.getBaseUrl();
+
+        hydraContainer = createHydraContainer()
+                .withNetworkAliases("hydra")
+                .withExposedPorts(4444, 4445)
+                .withEnv("DSN", DSN)
+                .withEnv("URLS_SELF_ISSUER", ISSUER)
+                .withEnv("URLS_CONSENT", loginAndConsentBaseUrl + "/consent")
+                .withEnv("URLS_LOGIN", loginAndConsentBaseUrl + "/login")
+                .withEnv("SERVE_TLS_KEY_PATH", "/tmp/certs/localhost.pem")
+                .withEnv("SERVE_TLS_CERT_PATH", "/tmp/certs/localhost.pem")
+                .withEnv("TTL_ACCESS_TOKEN", ttlAccessToken.getSeconds() + "s")
+                .withEnv("STRATEGIES_ACCESS_TOKEN", useJwt ? "jwt" : null)
+                .withEnv("LOG_LEAK_SENSITIVE_VALUES", "true")
+                .withCommand("serve", "all")
+                .withCopyFileToContainer(MountableFile.forClasspathResource("/cert"), "/tmp/certs")
+                .waitingFor(new WaitAllStrategy()
+                        .withStrategy(Wait.forLogMessage(".*Setting up http server on :4444.*", 1))
+                        .withStrategy(Wait.forLogMessage(".*Setting up http server on :4445.*", 1)));
+        if (exposeFixedPorts) {
+            hydraContainer = hydraContainer
+                    .withFixedExposedPort(4444, 4444)
+                    .withFixedExposedPort(4445, 4445);
+        }
+        closer.register(hydraContainer);
+        hydraContainer.start();
+    }
+
+    public FixedHostPortGenericContainer createHydraContainer()
+    {
+        return new FixedHostPortGenericContainer<>(HYDRA_IMAGE).withNetwork(network);
+    }
+
+    public void createClient(
+            String clientId,
+            String clientSecret,
+            TokenEndpointAuthMethod tokenEndpointAuthMethod,
+            List audiences,
+            String callbackUrl)
+    {
+        createHydraContainer()
+                .withCommand("clients", "create",
+                        "--endpoint", "https://hydra:4445",
+                        "--skip-tls-verify",
+                        "--id", clientId,
+                        "--secret", clientSecret,
+                        "--audience", String.join(",", audiences),
+                        "--grant-types", "authorization_code,refresh_token,client_credentials",
+                        "--response-types", "token,code,id_token",
+                        "--scope", "openid,offline",
+                        "--token-endpoint-auth-method", tokenEndpointAuthMethod.getValue(),
+                        "--callbacks", callbackUrl)
+                .withStartupCheckStrategy(new OneShotStartupCheckStrategy().withTimeout(Duration.ofSeconds(30)))
+                .start();
+    }
+
+    public String getToken(String clientId, String clientSecret, List audiences)
+            throws IOException
+    {
+        try (Response response = httpClient
+                .newCall(
+                        new Request.Builder()
+                                .url("https://localhost:" + getAuthPort() + "/oauth2/token")
+                                .addHeader(HttpHeaders.AUTHORIZATION, Credentials.basic(clientId, clientSecret))
+                                .post(new FormBody.Builder()
+                                        .add("grant_type", GrantType.CLIENT_CREDENTIALS.getValue())
+                                        .add("audience", String.join(" ", audiences))
+                                        .build())
+                                .build())
+                .execute()) {
+            checkState(response.code() == SC_OK);
+            requireNonNull(response.body());
+            return mapper.readTree(response.body().byteStream())
+                    .get("access_token")
+                    .textValue();
+        }
+    }
+
+    public int getAuthPort()
+    {
+        return hydraContainer.getMappedPort(4444);
+    }
+
+    public int getAdminPort()
+    {
+        return hydraContainer.getMappedPort(4445);
+    }
+
+    @Override
+    public void close()
+            throws IOException
+    {
+        try {
+            closer.close();
+        }
+        catch (Exception e) {
+            throwIfUnchecked(e);
+            throw new RuntimeException(e);
+        }
+    }
+
+    private TestingHttpServer createTestingLoginAndConsentServer()
+            throws IOException
+    {
+        NodeInfo nodeInfo = new NodeInfo("test");
+        HttpServerConfig config = new HttpServerConfig().setHttpPort(0);
+        HttpServerInfo httpServerInfo = new HttpServerInfo(config, nodeInfo);
+        return new TestingHttpServer(
+                httpServerInfo,
+                nodeInfo,
+                config,
+                new AcceptAllLoginsAndConsentsServlet(),
+                ImmutableMap.of(),
+                ImmutableMap.of(),
+                Optional.empty());
+    }
+
+    private class AcceptAllLoginsAndConsentsServlet
+            extends HttpServlet
+    {
+        private final ObjectMapper mapper = new ObjectMapper();
+        private final OkHttpClient httpClient;
+
+        public AcceptAllLoginsAndConsentsServlet()
+        {
+            OkHttpClient.Builder httpClientBuilder = new OkHttpClient.Builder();
+            setupInsecureSsl(httpClientBuilder);
+            httpClient = httpClientBuilder.build();
+        }
+
+        @Override
+        protected void doGet(HttpServletRequest request, HttpServletResponse response)
+                throws IOException
+        {
+            if (request.getPathInfo().equals("/login")) {
+                acceptLogin(request, response);
+                return;
+            }
+            if (request.getPathInfo().contains("/consent")) {
+                acceptConsent(request, response);
+                return;
+            }
+            response.setStatus(SC_NOT_FOUND);
+        }
+
+        private void acceptLogin(HttpServletRequest request, HttpServletResponse response)
+                throws IOException
+        {
+            String loginChallenge = request.getParameter("login_challenge");
+            try (Response loginAcceptResponse = acceptLogin(loginChallenge)) {
+                sendRedirect(loginAcceptResponse, response);
+            }
+        }
+
+        private void acceptConsent(HttpServletRequest request, HttpServletResponse response)
+                throws IOException
+        {
+            String consentChallenge = request.getParameter("consent_challenge");
+            JsonNode consentRequest = getConsentRequest(consentChallenge);
+            try (Response acceptConsentResponse = acceptConsent(consentChallenge, consentRequest)) {
+                sendRedirect(acceptConsentResponse, response);
+            }
+        }
+
+        private Response acceptLogin(String loginChallenge)
+                throws IOException
+        {
+            return httpClient.newCall(
+                            new Request.Builder()
+                                    .url("https://localhost:" + getAdminPort() + "/oauth2/auth/requests/login/accept?login_challenge=" + loginChallenge)
+                                    .put(RequestBody.create(
+                                            MediaType.parse(APPLICATION_JSON),
+                                            mapper.writeValueAsString(mapper.createObjectNode().put("subject", "foo@bar.com"))))
+                                    .build())
+                    .execute();
+        }
+
+        private JsonNode getConsentRequest(String consentChallenge)
+                throws IOException
+        {
+            try (Response response = httpClient.newCall(
+                            new Request.Builder()
+                                    .url("https://localhost:" + getAdminPort() + "/oauth2/auth/requests/consent?consent_challenge=" + consentChallenge)
+                                    .get()
+                                    .build())
+                    .execute()) {
+                requireNonNull(response.body());
+                return mapper.readTree(response.body().byteStream());
+            }
+        }
+
+        private Response acceptConsent(String consentChallenge, JsonNode consentRequest)
+                throws IOException
+        {
+            return httpClient.newCall(
+                            new Request.Builder()
+                                    .url("https://localhost:" + getAdminPort() + "/oauth2/auth/requests/consent/accept?consent_challenge=" + consentChallenge)
+                                    .put(RequestBody.create(
+                                            MediaType.parse(APPLICATION_JSON),
+                                            mapper.writeValueAsString(mapper.createObjectNode()
+                                                    .set("grant_scope", consentRequest.get("requested_scope"))
+                                                    .set("grant_access_token_audience", consentRequest.get("requested_access_token_audience")))))
+                                    .build())
+                    .execute();
+        }
+
+        private void sendRedirect(Response redirectResponse, HttpServletResponse response)
+                throws IOException
+        {
+            requireNonNull(redirectResponse.body());
+            response.sendRedirect(
+                    toHostUrl(mapper.readTree(redirectResponse.body().byteStream())
+                            .get("redirect_to")
+                            .textValue()));
+        }
+
+        private String toHostUrl(String url)
+        {
+            return HttpUrl.get(URI.create(url))
+                    .newBuilder()
+                    .port(getAuthPort())
+                    .toString();
+        }
+    }
+
+    private static void runTestServer(boolean useJwt)
+            throws Exception
+    {
+        try (TestingHydraIdentityProvider service = new TestingHydraIdentityProvider(Duration.ofMinutes(30), useJwt, true)) {
+            service.start();
+            service.createClient(
+                    "presto-client",
+                    "presto-secret",
+                    CLIENT_SECRET_BASIC,
+                    ImmutableList.of("https://localhost:8443/ui"),
+                    "https://localhost:8443/oauth2/callback");
+            ImmutableMap.Builder config = ImmutableMap.builder()
+                    .put("http-server.https.port", "8443")
+                    .put("http-server.https.enabled", "true")
+                    .put("http-server.https.keystore.path", Resources.getResource("cert/localhost.pem").getPath())
+                    .put("http-server.https.keystore.key", "")
+                    .put("http-server.authentication.type", "OAUTH2")
+                    .put("http-server.authentication.oauth2.issuer", ISSUER)
+                    .put("http-server.authentication.oauth2.client-id", "presto-client")
+                    .put("http-server.authentication.oauth2.client-secret", "presto-secret")
+                    .put("http-server.authentication.oauth2.user-mapping.pattern", "(.*)@.*")
+                    .put("http-server.authentication.oauth2.oidc.use-userinfo-endpoint", String.valueOf(!useJwt))
+                    .put("oauth2-jwk.http-client.trust-store-path", Resources.getResource("cert/localhost.pem").getPath());
+            try (TestingPrestoServer server = new TestingPrestoServer(config.build())) {
+                server.getInstance(Key.get(OAuth2Client.class)).load();
+                Thread.sleep(Long.MAX_VALUE);
+            }
+        }
+    }
+
+    public static void main(String[] args)
+            throws Exception
+    {
+        Logging logging = Logging.initialize();
+        try {
+            logging.setLevel(OAuth2Service.class.getName(), Level.DEBUG);
+            runTestServer(false);
+        }
+        finally {
+            logging.setLevel(OAuth2Service.class.getName(), Level.INFO);
+        }
+    }
+}
diff --git a/presto-main/src/test/java/com/facebook/presto/server/security/oauth2/TokenEndpointAuthMethod.java b/presto-main/src/test/java/com/facebook/presto/server/security/oauth2/TokenEndpointAuthMethod.java
new file mode 100644
index 0000000000000..c05e636e9b11c
--- /dev/null
+++ b/presto-main/src/test/java/com/facebook/presto/server/security/oauth2/TokenEndpointAuthMethod.java
@@ -0,0 +1,33 @@
+/*
+ * Licensed 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 com.facebook.presto.server.security.oauth2;
+
+import static java.util.Objects.requireNonNull;
+
+public enum TokenEndpointAuthMethod
+{
+    CLIENT_SECRET_BASIC("client_secret_basic");
+
+    private final String value;
+
+    TokenEndpointAuthMethod(String value)
+    {
+        this.value = requireNonNull(value, "value is null");
+    }
+
+    public String getValue()
+    {
+        return value;
+    }
+}
diff --git a/presto-main/src/test/resources/cert/generate.sh b/presto-main/src/test/resources/cert/generate.sh
new file mode 100755
index 0000000000000..a84749045a5e4
--- /dev/null
+++ b/presto-main/src/test/resources/cert/generate.sh
@@ -0,0 +1,7 @@
+#!/bin/sh
+
+set -eux
+
+openssl req -new -x509 -newkey rsa:4096 -sha256 -nodes -keyout localhost.key -days 3560 -out localhost.crt -config localhost.conf
+cat localhost.crt localhost.key > localhost.pem
+
diff --git a/presto-main/src/test/resources/cert/localhost.conf b/presto-main/src/test/resources/cert/localhost.conf
new file mode 100644
index 0000000000000..560e1a454c8d9
--- /dev/null
+++ b/presto-main/src/test/resources/cert/localhost.conf
@@ -0,0 +1,20 @@
+[req]
+default_bits = 4096
+prompt = no
+default_md = sha256
+x509_extensions = v3_req
+distinguished_name = dn
+
+[dn]
+C = US
+ST = California
+L = Palo Alto
+O = PrestoTest
+CN = PrestoTest
+
+[v3_req]
+subjectAltName = @alt_names
+
+[alt_names]
+IP.1   = 127.0.0.1
+DNS.1   = localhost
diff --git a/presto-main/src/test/resources/cert/localhost.pem b/presto-main/src/test/resources/cert/localhost.pem
new file mode 100644
index 0000000000000..2d8a8fb289876
--- /dev/null
+++ b/presto-main/src/test/resources/cert/localhost.pem
@@ -0,0 +1,83 @@
+-----BEGIN CERTIFICATE-----
+MIIFYTCCA0mgAwIBAgIJAKUofzuCtcfnMA0GCSqGSIb3DQEBCwUAMGAxCzAJBgNV
+BAYTAlVTMRMwEQYDVQQIDApDYWxpZm9ybmlhMRIwEAYDVQQHDAlQYWxvIEFsdG8x
+EzARBgNVBAoMClByZXN0b1Rlc3QxEzARBgNVBAMMClByZXN0b1Rlc3QwHhcNMjIw
+OTA2MTgyOTUxWhcNMzIwNjA1MTgyOTUxWjBgMQswCQYDVQQGEwJVUzETMBEGA1UE
+CAwKQ2FsaWZvcm5pYTESMBAGA1UEBwwJUGFsbyBBbHRvMRMwEQYDVQQKDApQcmVz
+dG9UZXN0MRMwEQYDVQQDDApQcmVzdG9UZXN0MIICIjANBgkqhkiG9w0BAQEFAAOC
+Ag8AMIICCgKCAgEA28P/OFPTMWu5AUt2YF3IvfFtZ2FRioB02+FIE21KtYM1r8w2
+2GRyvLuT8vBaYoh8bNSGI2x2R1NtfOalaUCUfr9XRyZPcP4FEE5x0QRK2SYYOfzr
+URx5gv3SSlhahmSjsFAojpG7lUBsKpopFcjZb0wSq3hFDVHQN57Xzmt1YHbTZrEt
+5yFyqt2AYRVHz8XxJbsUOy514/YGQfLLZqukSLYk055qFIclzFqXU+/cg6UVpl7U
+hlLTo0GApBQ2eLGCBDZqXhkCf2U1lMGGVLsJFNmGaumLV88yZmRYQC9MJWlCPCGG
+ZcyKNxjZq70SbmjIA6s0FVcXYZ0z6xQqDpVBichLebrtR8ShKU29u1ITL2kaF1iK
+gZi8FzEKwnzxLlTZACeBfGeELl9HKUmUvOwU9LHp0UX4fZLlcj5IYXk0IMExFEYa
+qfKYxThmdo1Gmpl1orW3mnx+BK8/VtMn7RquHTgQof/dry3EAMo+7/N9v3tI4m+5
+99g+DoHyOdpWYOsTtGmkYcWLG8/ka2lzcaLx347VBRBgNa/afKquPK3ogi04wYYl
+K8wBciyhg+J7MkC0k1Q5ek25qynotHLkIGC+LsmEOuw8kApc8/ZFvA5yavn4Dsm5
+x94CmVhJQqL5Dqr6IEmzxvWjFWCC4mD2os52Ff1RlEGuEB5B9CAeaWCSuGECAwEA
+AaMeMBwwGgYDVR0RBBMwEYcEfwAAAYIJbG9jYWxob3N0MA0GCSqGSIb3DQEBCwUA
+A4ICAQCnGqgmdRRIuEij5bAHHF+CkbtkHbOO0fTxmd5/Cd5vkiCUR0uhwBg9JUbC
+5zl/wBTx8guc0b7Wjvr3gwBUWQhGR2LKQMfatFZSy8nPcfSkZUGjY4s7KYI8mjPY
+1Lri/EE1gu2p+iB9Dw4EnHW3QSneyy4yrkpLcjywbkF93SsThgHQ27gEK3/HeGzY
+dFQY17z6zRNnlkMU2JVh0VptE6xtdR3WAMhscVx4dpEjFz9FSKnAYE2svTAy0OTD
+8+W795a4/eVCxdHw/3PqR4XSO8isdSimeSTtdpRqsDrW4jx4IIGVFMQPin5XABAl
+Wbbs8VMZPB1OvLSpmPtV79o/EklWS9x0MtbXF6iT5VBIbP1JoQLHQ9O0+V4PT0CH
+8f8+Wc2UqLskqJZOyDzV3Y81mcdmKYcWNN8LJBk6PA1Pt/RDAg58QW1KfUVleey8
+eaAMw7d9BoIM/nUe0Q4TPll0WJnFMWHLUItgs7YuN/39kOnxGQi3iKQCMV2qwRip
+tHqTvw3fHXQxEbZZhLxInC04+pOF3ZpqdzkeaZsXpwklV+uDw1u+5rzlheVrErwA
+BzAXENpJKZLs4mjHko+Z4loNzGsJzRci/YRkwSKW6stW04RWWEKQkc9YhFsDYDmt
+I70j8+R7et3nx2DqaQ18WgQ/5xVsEoXg5sgFjj1cRVFDQdrexA==
+-----END CERTIFICATE-----
+-----BEGIN PRIVATE KEY-----
+MIIJQgIBADANBgkqhkiG9w0BAQEFAASCCSwwggkoAgEAAoICAQDbw/84U9Mxa7kB
+S3ZgXci98W1nYVGKgHTb4UgTbUq1gzWvzDbYZHK8u5Py8FpiiHxs1IYjbHZHU218
+5qVpQJR+v1dHJk9w/gUQTnHRBErZJhg5/OtRHHmC/dJKWFqGZKOwUCiOkbuVQGwq
+mikVyNlvTBKreEUNUdA3ntfOa3VgdtNmsS3nIXKq3YBhFUfPxfEluxQ7LnXj9gZB
+8stmq6RItiTTnmoUhyXMWpdT79yDpRWmXtSGUtOjQYCkFDZ4sYIENmpeGQJ/ZTWU
+wYZUuwkU2YZq6YtXzzJmZFhAL0wlaUI8IYZlzIo3GNmrvRJuaMgDqzQVVxdhnTPr
+FCoOlUGJyEt5uu1HxKEpTb27UhMvaRoXWIqBmLwXMQrCfPEuVNkAJ4F8Z4QuX0cp
+SZS87BT0senRRfh9kuVyPkhheTQgwTEURhqp8pjFOGZ2jUaamXWitbeafH4Erz9W
+0yftGq4dOBCh/92vLcQAyj7v832/e0jib7n32D4OgfI52lZg6xO0aaRhxYsbz+Rr
+aXNxovHfjtUFEGA1r9p8qq48reiCLTjBhiUrzAFyLKGD4nsyQLSTVDl6TbmrKei0
+cuQgYL4uyYQ67DyQClzz9kW8DnJq+fgOybnH3gKZWElCovkOqvogSbPG9aMVYILi
+YPaiznYV/VGUQa4QHkH0IB5pYJK4YQIDAQABAoICAFSVcz2yxa5Xz7T33m/oqMKy
+kXEgu8ma919JrfwMLJ0AC0HGT7Wps5+AcskmSSNzdLBOe/JWZI+/RHy2KSQBfyXp
+byYrUJgkrL5B8vyHsmcxilGHTurBEuOf3bhPmUfwpC/QKkv1O0WOrhMXkoiX7Vgw
+516nw6wEuScvM9B2+45NLcBwoUI8VW3+ItM65ZDKlq32+ypsD2PV5UKsuCykE28I
+69OnPRz5h0rH80aTI0Rn3ZVTGmk4p8xGAcUlInIBoBEPAJGG/rcZtS2z7ofeFPi5
+YEr16HO7g6J1LKJHkf80LBIItTmpJ+lc3yqCcv2bxp/i3QD5rD4dy0XHVQiX2cj/
+KQSh64MAszRtXZE+GOIAT6UOxyEZqYrPNopXZWEROpM4W/c20HsA1ffpXu0abOcY
+rkZhWmKk5mJNBAe+EE6aoxGJ9Th6x7WO7WSjiAmXIxjoAk0YClOwMf6y8+5fuqg8
+aEtG4GOz2bfUksbE3AzaXMcu2o4tt+4+paZA0pkhB9Za1ySP5N4pCQ1kvKhdstZi
+e9laJtBS3YSS+rm2pe3g/rA9x5OrTDpms8S3LvFflG1KdPzJeBX7zRg6AHFBhAPI
+OpFvsdeVDHOmGe/MYjC0gpLkoVBJLY7cqiicqUUWYaFm0JMoma8VS2h+Y4+cV39N
+erXXLflO+Zn1EhiWrWYBAoIBAQD8K3evijZepPil3wtLff/dHgWzOsIubkBWEALu
+M+Ab9zm0ahZ1muEej3p0U+Eq8TUbOAqaOWWndFMiqTcjco3zluSIh7wYajY827uG
+5Fuyw3zIUm1g5HicZyVsqBJxRcSxkNqzFMMnvhFqUNyrsKHL+/GsCnWiuZhz8z8h
+j9l616DF6LpefZrDL1nne4xfliZuE99BWNBi7n1n48q/yxpZ+PyNjAZFUu2wcOFE
++/z5wuuNwywk0gsWEPiqt1LAIbLwHYaTfIwK/OO19ScfFjzw8ZRNLcNEhFbFyZfR
+aVkL22uqU7vtBzA1ld2AtogDTEdMQ1ruFd/ZIoOiHK6OpUhJAoIBAQDfGoeTi51A
+Fh6gDh4/eiXKpggUvxBbLOkyfcb3j9VCbPmm+09YYR+b99FD7jhM2gKY1QXscvm5
+LFwMiXBNOv110PL1H0lGaNnv5sRFgQwAd7BpBNSAxYQCjF7RZz1GBHq3iLaAYtY4
+jRVPGL+n55E84jzi2Ip856Nf4TaTf98rsslSPbr6zwFvsGALFOki7hJivcelS1Qm
+3U6UuJ3CmWv6GA3BY6BCj+5pXzR8+5bisYuIFUm1uW6sSn+1aaogWtOTRIFqf3mT
+WSewoOcJpXHs5kBHkRVMycTIUaijbMdoANQzsZ8d6dfcyxpfrncxEaf+VJWKxEVp
+OVSTUbpn/d9ZAoIBAEBF0/up/rGg2r9sWjSjqNNzE4DbOSMcdsl4y0ZrcnOuT8bs
+Q002bKqdZ1i/CGUplZ+aaRlmB8Lmo0nyV1txlzy++QDTl92hNLHOT73R9o1ZxjRI
+zhgkI5m5sJBBRnIYlkmr4hJC+HrotweiFJyuKI8VaEOxZspTA8iJ901WnNfyncfT
+yazL1uZo60FU/DJg0uq7pevB91s/7jbMmKDJ462LCNQLHI4O1QZjvwcWMyR1yhQX
+6uh3oNu+96KLl0vhSvpojCSLWiZyzpdSJOaHhIDlEieZwmt0T6mZ+FgnwcqD4q1H
+Kl7/tgnyaMKlw4UTrBiEEmkcqjFt2p83MEarWgECggEBAKzFnrhkJiK6/nx0cng1
+345PhXKLg98XqH+xZ6PPfxcxzSPC+m82x4PBJg21LWRWcCxqy2uQnlMIR0BuLsmg
+JShX585rrBMan6toyhYJGYJDLhol42rViqVujv8bNBhE38PB25MQ91RT7WyTfdhJ
+O/AqQ3xotNaFi790aQ9Qt0Lf8Yf+xg30wOf9bmMmjmS+eP5+eV1IOKLgPzpsvb81
+kKjcd8qLnE/vpnFziPJA41gqpiN8WNiiAVLrXnremSD1NWOWaaJPlZbGNDZUZJbT
+yKXsqVrCv/v3RKzcj/v/AW1JNwvRQaeor8IMhyARu7wEMFSErEoKNLaH7zcm03Q0
+5gECggEAcaeUlvCqcQiU/i42f9cHu2wx7kDKil6JS8ZYsJ0dGl5g9u2+VyeUBJvI
+jLTrnwaGyn+Gx+PlGaC0QK5i9WjbblxrPuWFJk63pGDzLW35jrM55QC9kiI4C1Q7
+TyWOfjGBGa+Gi+v5bo+B3osgAHx7F2SQaoZ2C6Qbo+CMvDCihYFsPcOVngNzf+7+
+QkPOOO5ixlopHyekKe6dcMvb1PzovSEK3DZoh0XH4M5hz3LNk3Z57seNyN9pxcLI
+U4qSYvcgYwbYZWjXymU6LVsx5lKJ4RrJGS56K0Q3ebC9nz7pNoVGKxvFAo4OI179
+pTI3VZB66Rqurj0RWq452LJE1Onb5A==
+-----END PRIVATE KEY-----
diff --git a/presto-main/src/test/resources/oidc/openid-configuration-invalid-issuer.json b/presto-main/src/test/resources/oidc/openid-configuration-invalid-issuer.json
new file mode 100644
index 0000000000000..2f23e8ac98154
--- /dev/null
+++ b/presto-main/src/test/resources/oidc/openid-configuration-invalid-issuer.json
@@ -0,0 +1,106 @@
+{
+    "issuer": "https://invalid-issuer.com",
+    "authorization_endpoint": "https://invalid-issuer.com/connect/authorize",
+    "token_endpoint": "https://invalid-issuer.com/connect/token",
+    "token_endpoint_auth_methods_supported": [
+        "client_secret_basic",
+        "private_key_jwt"
+    ],
+    "token_endpoint_auth_signing_alg_values_supported": [
+        "RS256",
+        "ES256"
+    ],
+    "userinfo_endpoint": "https://invalid-issuer.com/connect/userinfo",
+    "check_session_iframe": "https://invalid-issuer.com/connect/check_session",
+    "end_session_endpoint": "https://invalid-issuer.com/connect/end_session",
+    "jwks_uri": "https://invalid-issuer.com/jwks.json",
+    "registration_endpoint": "https://invalid-issuer.com/connect/register",
+    "scopes_supported": [
+        "openid",
+        "profile",
+        "email",
+        "address",
+        "phone",
+        "offline_access"
+    ],
+    "response_types_supported": [
+        "code",
+        "code id_token",
+        "id_token",
+        "token id_token"
+    ],
+    "acr_values_supported": [
+        "urn:mace:incommon:iap:silver",
+        "urn:mace:incommon:iap:bronze"
+    ],
+    "subject_types_supported": [
+        "public",
+        "pairwise"
+    ],
+    "userinfo_signing_alg_values_supported": [
+        "RS256",
+        "ES256",
+        "HS256"
+    ],
+    "userinfo_encryption_alg_values_supported": [
+        "RSA1_5",
+        "A128KW"
+    ],
+    "userinfo_encryption_enc_values_supported": [
+        "A128CBC-HS256",
+        "A128GCM"
+    ],
+    "id_token_signing_alg_values_supported": [
+        "RS256",
+        "ES256",
+        "HS256"
+    ],
+    "id_token_encryption_alg_values_supported": [
+        "RSA1_5",
+        "A128KW"
+    ],
+    "id_token_encryption_enc_values_supported": [
+        "A128CBC-HS256",
+        "A128GCM"
+    ],
+    "request_object_signing_alg_values_supported": [
+        "none",
+        "RS256",
+        "ES256"
+    ],
+    "display_values_supported": [
+        "page",
+        "popup"
+    ],
+    "claim_types_supported": [
+        "normal",
+        "distributed"
+    ],
+    "claims_supported": [
+        "sub",
+        "iss",
+        "auth_time",
+        "acr",
+        "name",
+        "given_name",
+        "family_name",
+        "nickname",
+        "profile",
+        "picture",
+        "website",
+        "email",
+        "email_verified",
+        "locale",
+        "zoneinfo",
+        "http://example.info/claims/groups"
+    ],
+    "claims_parameter_supported": true,
+    "service_documentation": "http://invalid-issuer.com/connect/service_documentation.html",
+    "ui_locales_supported": [
+        "en-US",
+        "en-GB",
+        "en-CA",
+        "fr-FR",
+        "fr-CA"
+    ]
+}
diff --git a/presto-main/src/test/resources/oidc/openid-configuration-with-access-token-issuer.json b/presto-main/src/test/resources/oidc/openid-configuration-with-access-token-issuer.json
new file mode 100644
index 0000000000000..905032ce3df1d
--- /dev/null
+++ b/presto-main/src/test/resources/oidc/openid-configuration-with-access-token-issuer.json
@@ -0,0 +1,87 @@
+{
+    "issuer": "https://issuer.com",
+    "authorization_endpoint": "https://issuer.com/connect/authorize",
+    "token_endpoint": "https://issuer.com/connect/token",
+    "jwks_uri": "https://issuer.com/jwks.json",
+    "token_endpoint_auth_methods_supported": [
+        "client_secret_post",
+        "client_secret_basic",
+        "private_key_jwt",
+        "windows_client_authentication"
+    ],
+    "response_types_supported": [
+        "code",
+        "id_token",
+        "code id_token",
+        "id_token token",
+        "code token",
+        "code id_token token"
+    ],
+    "response_modes_supported": [
+        "query",
+        "fragment",
+        "form_post"
+    ],
+    "grant_types_supported": [
+        "authorization_code",
+        "refresh_token",
+        "client_credentials",
+        "urn:ietf:params:oauth:grant-type:jwt-bearer",
+        "implicit",
+        "password",
+        "srv_challenge",
+        "urn:ietf:params:oauth:grant-type:device_code",
+        "device_code"
+    ],
+    "subject_types_supported": [
+        "pairwise"
+    ],
+    "scopes_supported": [
+        "user_impersonation",
+        "vpn_cert",
+        "email",
+        "openid",
+        "profile",
+        "allatclaims",
+        "logon_cert",
+        "aza",
+        "winhello_cert"
+    ],
+    "id_token_signing_alg_values_supported": [
+        "RS256"
+    ],
+    "token_endpoint_auth_signing_alg_values_supported": [
+        "RS256"
+    ],
+    "access_token_issuer": "http://access-token-issuer.com/adfs/services/trust",
+    "claims_supported": [
+        "aud",
+        "iss",
+        "iat",
+        "exp",
+        "auth_time",
+        "nonce",
+        "at_hash",
+        "c_hash",
+        "sub",
+        "upn",
+        "unique_name",
+        "pwd_url",
+        "pwd_exp",
+        "mfa_auth_time",
+        "sid",
+        "nbf"
+    ],
+    "microsoft_multi_refresh_token": true,
+    "userinfo_endpoint": "https://issuer.com/connect/userinfo",
+    "capabilities": [],
+    "end_session_endpoint": "https://issuer.com/adfs/oauth2/logout",
+    "as_access_token_token_binding_supported": true,
+    "as_refresh_token_token_binding_supported": true,
+    "resource_access_token_token_binding_supported": true,
+    "op_id_token_token_binding_supported": true,
+    "rp_id_token_token_binding_supported": true,
+    "frontchannel_logout_supported": true,
+    "frontchannel_logout_session_supported": true,
+    "device_authorization_endpoint": "https://issuer.com/adfs/oauth2/devicecode"
+}
diff --git a/presto-main/src/test/resources/oidc/openid-configuration-without-userinfo.json b/presto-main/src/test/resources/oidc/openid-configuration-without-userinfo.json
new file mode 100644
index 0000000000000..96f8c6048acf6
--- /dev/null
+++ b/presto-main/src/test/resources/oidc/openid-configuration-without-userinfo.json
@@ -0,0 +1,92 @@
+{
+    "issuer": "https://issuer.com",
+    "authorization_endpoint": "https://issuer.com/connect/authorize",
+    "token_endpoint": "https://issuer.com/connect/token",
+    "token_endpoint_auth_methods_supported": [
+        "client_secret_basic",
+        "private_key_jwt"
+    ],
+    "token_endpoint_auth_signing_alg_values_supported": [
+        "RS256",
+        "ES256"
+    ],
+    "check_session_iframe": "https://issuer.com/connect/check_session",
+    "end_session_endpoint": "https://issuer.com/connect/end_session",
+    "jwks_uri": "https://issuer.com/jwks.json",
+    "registration_endpoint": "https://issuer.com/connect/register",
+    "scopes_supported": [
+        "openid",
+        "profile",
+        "email",
+        "address",
+        "phone",
+        "offline_access"
+    ],
+    "response_types_supported": [
+        "code",
+        "code id_token",
+        "id_token",
+        "token id_token"
+    ],
+    "acr_values_supported": [
+        "urn:mace:incommon:iap:silver",
+        "urn:mace:incommon:iap:bronze"
+    ],
+    "subject_types_supported": [
+        "public",
+        "pairwise"
+    ],
+    "id_token_signing_alg_values_supported": [
+        "RS256",
+        "ES256",
+        "HS256"
+    ],
+    "id_token_encryption_alg_values_supported": [
+        "RSA1_5",
+        "A128KW"
+    ],
+    "id_token_encryption_enc_values_supported": [
+        "A128CBC-HS256",
+        "A128GCM"
+    ],
+    "request_object_signing_alg_values_supported": [
+        "none",
+        "RS256",
+        "ES256"
+    ],
+    "display_values_supported": [
+        "page",
+        "popup"
+    ],
+    "claim_types_supported": [
+        "normal",
+        "distributed"
+    ],
+    "claims_supported": [
+        "sub",
+        "iss",
+        "auth_time",
+        "acr",
+        "name",
+        "given_name",
+        "family_name",
+        "nickname",
+        "profile",
+        "picture",
+        "website",
+        "email",
+        "email_verified",
+        "locale",
+        "zoneinfo",
+        "http://example.info/claims/groups"
+    ],
+    "claims_parameter_supported": true,
+    "service_documentation": "http://issuer.com/connect/service_documentation.html",
+    "ui_locales_supported": [
+        "en-US",
+        "en-GB",
+        "en-CA",
+        "fr-FR",
+        "fr-CA"
+    ]
+}
diff --git a/presto-main/src/test/resources/oidc/openid-configuration.json b/presto-main/src/test/resources/oidc/openid-configuration.json
new file mode 100644
index 0000000000000..39c456fca85b6
--- /dev/null
+++ b/presto-main/src/test/resources/oidc/openid-configuration.json
@@ -0,0 +1,106 @@
+{
+    "issuer": "https://issuer.com",
+    "authorization_endpoint": "https://issuer.com/connect/authorize",
+    "token_endpoint": "https://issuer.com/connect/token",
+    "token_endpoint_auth_methods_supported": [
+        "client_secret_basic",
+        "private_key_jwt"
+    ],
+    "token_endpoint_auth_signing_alg_values_supported": [
+        "RS256",
+        "ES256"
+    ],
+    "userinfo_endpoint": "https://issuer.com/connect/userinfo",
+    "check_session_iframe": "https://issuer.com/connect/check_session",
+    "end_session_endpoint": "https://issuer.com/connect/end_session",
+    "jwks_uri": "https://issuer.com/jwks.json",
+    "registration_endpoint": "https://issuer.com/connect/register",
+    "scopes_supported": [
+        "openid",
+        "profile",
+        "email",
+        "address",
+        "phone",
+        "offline_access"
+    ],
+    "response_types_supported": [
+        "code",
+        "code id_token",
+        "id_token",
+        "token id_token"
+    ],
+    "acr_values_supported": [
+        "urn:mace:incommon:iap:silver",
+        "urn:mace:incommon:iap:bronze"
+    ],
+    "subject_types_supported": [
+        "public",
+        "pairwise"
+    ],
+    "userinfo_signing_alg_values_supported": [
+        "RS256",
+        "ES256",
+        "HS256"
+    ],
+    "userinfo_encryption_alg_values_supported": [
+        "RSA1_5",
+        "A128KW"
+    ],
+    "userinfo_encryption_enc_values_supported": [
+        "A128CBC-HS256",
+        "A128GCM"
+    ],
+    "id_token_signing_alg_values_supported": [
+        "RS256",
+        "ES256",
+        "HS256"
+    ],
+    "id_token_encryption_alg_values_supported": [
+        "RSA1_5",
+        "A128KW"
+    ],
+    "id_token_encryption_enc_values_supported": [
+        "A128CBC-HS256",
+        "A128GCM"
+    ],
+    "request_object_signing_alg_values_supported": [
+        "none",
+        "RS256",
+        "ES256"
+    ],
+    "display_values_supported": [
+        "page",
+        "popup"
+    ],
+    "claim_types_supported": [
+        "normal",
+        "distributed"
+    ],
+    "claims_supported": [
+        "sub",
+        "iss",
+        "auth_time",
+        "acr",
+        "name",
+        "given_name",
+        "family_name",
+        "nickname",
+        "profile",
+        "picture",
+        "website",
+        "email",
+        "email_verified",
+        "locale",
+        "zoneinfo",
+        "http://example.info/claims/groups"
+    ],
+    "claims_parameter_supported": true,
+    "service_documentation": "http://issuer.com/connect/service_documentation.html",
+    "ui_locales_supported": [
+        "en-US",
+        "en-GB",
+        "en-CA",
+        "fr-FR",
+        "fr-CA"
+    ]
+}
diff --git a/presto-memory/src/main/java/com/facebook/presto/plugin/memory/MemoryMetadata.java b/presto-memory/src/main/java/com/facebook/presto/plugin/memory/MemoryMetadata.java
index 82232eb153f70..b9741caf478c9 100644
--- a/presto-memory/src/main/java/com/facebook/presto/plugin/memory/MemoryMetadata.java
+++ b/presto-memory/src/main/java/com/facebook/presto/plugin/memory/MemoryMetadata.java
@@ -356,7 +356,7 @@ private void updateRowsOnHosts(MemoryTableHandle table, Collection fragme
     }
 
     @Override
-    public synchronized List getTableLayouts(
+    public synchronized ConnectorTableLayoutResult getTableLayoutForConstraint(
             ConnectorSession session,
             ConnectorTableHandle handle,
             Constraint constraint,
@@ -375,7 +375,7 @@ public synchronized List getTableLayouts(
                 tableDataFragments.get(memoryTableHandle.getTableId()).values());
 
         MemoryTableLayoutHandle layoutHandle = new MemoryTableLayoutHandle(memoryTableHandle, expectedFragments);
-        return ImmutableList.of(new ConnectorTableLayoutResult(getTableLayout(session, layoutHandle), constraint.getSummary()));
+        return new ConnectorTableLayoutResult(getTableLayout(session, layoutHandle), constraint.getSummary());
     }
 
     @Override
diff --git a/presto-memory/src/test/java/com/facebook/presto/plugin/memory/TestMemoryMetadata.java b/presto-memory/src/test/java/com/facebook/presto/plugin/memory/TestMemoryMetadata.java
index 529b46cf580e4..3eb2834008551 100644
--- a/presto-memory/src/test/java/com/facebook/presto/plugin/memory/TestMemoryMetadata.java
+++ b/presto-memory/src/test/java/com/facebook/presto/plugin/memory/TestMemoryMetadata.java
@@ -150,9 +150,9 @@ public void testReadTableBeforeCreationCompleted()
         assertTrue(tableNames.size() == 1, "Expected exactly one table");
 
         ConnectorTableHandle tableHandle = metadata.getTableHandle(SESSION, tableName);
-        List tableLayouts = metadata.getTableLayouts(SESSION, tableHandle, Constraint.alwaysTrue(), Optional.empty());
-        assertTrue(tableLayouts.size() == 1, "Expected exactly one layout.");
-        ConnectorTableLayout tableLayout = tableLayouts.get(0).getTableLayout();
+        ConnectorTableLayoutResult tableLayoutResult = metadata.getTableLayoutForConstraint(SESSION, tableHandle, Constraint.alwaysTrue(), Optional.empty());
+        assertTrue(tableLayoutResult != null, "Table layout is null.");
+        ConnectorTableLayout tableLayout = tableLayoutResult.getTableLayout();
         ConnectorTableLayoutHandle tableLayoutHandle = tableLayout.getHandle();
         assertTrue(tableLayoutHandle instanceof MemoryTableLayoutHandle);
         assertTrue(((MemoryTableLayoutHandle) tableLayoutHandle).getDataFragments().isEmpty(), "Data fragments should be empty");
diff --git a/presto-ml/src/test/java/com/facebook/presto/ml/AbstractTestMLFunctions.java b/presto-ml/src/test/java/com/facebook/presto/ml/AbstractTestMLFunctions.java
index dffbefd9189aa..4ee697957e540 100644
--- a/presto-ml/src/test/java/com/facebook/presto/ml/AbstractTestMLFunctions.java
+++ b/presto-ml/src/test/java/com/facebook/presto/ml/AbstractTestMLFunctions.java
@@ -15,13 +15,21 @@
 package com.facebook.presto.ml;
 
 import com.facebook.presto.operator.scalar.AbstractTestFunctions;
+import com.facebook.presto.sql.analyzer.FeaturesConfig;
+import com.facebook.presto.sql.analyzer.FunctionsConfig;
 import org.testng.annotations.BeforeClass;
 
+import static com.facebook.presto.SessionTestUtils.TEST_SESSION;
 import static com.facebook.presto.metadata.FunctionExtractor.extractFunctions;
 
 abstract class AbstractTestMLFunctions
         extends AbstractTestFunctions
 {
+    public AbstractTestMLFunctions()
+    {
+        super(TEST_SESSION, new FeaturesConfig(), new FunctionsConfig(), false);
+    }
+
     @BeforeClass
     protected void registerFunctions()
     {
diff --git a/presto-mongodb/pom.xml b/presto-mongodb/pom.xml
index e8a6429ea6b4e..8b4d7b916f61e 100644
--- a/presto-mongodb/pom.xml
+++ b/presto-mongodb/pom.xml
@@ -15,7 +15,7 @@
     
         ${project.parent.basedir}
         3.12.14
-        1.5.0
+        1.47.0
         true
     
 
@@ -222,17 +222,15 @@
                 
             
         
+        
         
-            test-mongo-distributed-queries
+            ci-full-tests
             
                 
                     
                         org.apache.maven.plugins
                         maven-surefire-plugin
                         
-                            
-                                **/TestMongoDistributedQueries.java
-                            
                         
                     
                 
diff --git a/presto-mongodb/src/main/java/com/facebook/presto/mongodb/MongoClientConfig.java b/presto-mongodb/src/main/java/com/facebook/presto/mongodb/MongoClientConfig.java
index 07987c290c365..b23570fdad83d 100644
--- a/presto-mongodb/src/main/java/com/facebook/presto/mongodb/MongoClientConfig.java
+++ b/presto-mongodb/src/main/java/com/facebook/presto/mongodb/MongoClientConfig.java
@@ -59,6 +59,7 @@ public class MongoClientConfig
     private WriteConcernType writeConcern = WriteConcernType.ACKNOWLEDGED;
     private String requiredReplicaSetName;
     private String implicitRowFieldPrefix = "_pos";
+    private boolean caseSensitiveNameMatchingEnabled;
 
     @NotNull
     public String getSchemaCollection()
@@ -328,4 +329,16 @@ public MongoClientConfig setSslEnabled(boolean sslEnabled)
         this.sslEnabled = sslEnabled;
         return this;
     }
+
+    public boolean isCaseSensitiveNameMatchingEnabled()
+    {
+        return caseSensitiveNameMatchingEnabled;
+    }
+
+    @Config("case-sensitive-name-matching")
+    public MongoClientConfig setCaseSensitiveNameMatchingEnabled(boolean caseSensitiveNameMatchingEnabled)
+    {
+        this.caseSensitiveNameMatchingEnabled = caseSensitiveNameMatchingEnabled;
+        return this;
+    }
 }
diff --git a/presto-mongodb/src/main/java/com/facebook/presto/mongodb/MongoConnector.java b/presto-mongodb/src/main/java/com/facebook/presto/mongodb/MongoConnector.java
index b21f6fdf4511c..f7feee5e7eaf3 100644
--- a/presto-mongodb/src/main/java/com/facebook/presto/mongodb/MongoConnector.java
+++ b/presto-mongodb/src/main/java/com/facebook/presto/mongodb/MongoConnector.java
@@ -41,18 +41,21 @@ public class MongoConnector
     private final MongoPageSinkProvider pageSinkProvider;
 
     private final ConcurrentMap transactions = new ConcurrentHashMap<>();
+    private final MongoClientConfig mongoClientConfig;
 
     @Inject
     public MongoConnector(
             MongoSession mongoSession,
             MongoSplitManager splitManager,
             MongoPageSourceProvider pageSourceProvider,
-            MongoPageSinkProvider pageSinkProvider)
+            MongoPageSinkProvider pageSinkProvider,
+            MongoClientConfig mongoClientConfig)
     {
         this.mongoSession = mongoSession;
         this.splitManager = requireNonNull(splitManager, "splitManager is null");
         this.pageSourceProvider = requireNonNull(pageSourceProvider, "pageSourceProvider is null");
         this.pageSinkProvider = requireNonNull(pageSinkProvider, "pageSinkProvider is null");
+        this.mongoClientConfig = requireNonNull(mongoClientConfig, "mongoClientConfig is null");
     }
 
     @Override
@@ -60,7 +63,7 @@ public ConnectorTransactionHandle beginTransaction(IsolationLevel isolationLevel
     {
         checkConnectorSupports(READ_UNCOMMITTED, isolationLevel);
         MongoTransactionHandle transaction = new MongoTransactionHandle();
-        transactions.put(transaction, new MongoMetadata(mongoSession));
+        transactions.put(transaction, new MongoMetadata(mongoSession, mongoClientConfig));
         return transaction;
     }
 
diff --git a/presto-mongodb/src/main/java/com/facebook/presto/mongodb/MongoMetadata.java b/presto-mongodb/src/main/java/com/facebook/presto/mongodb/MongoMetadata.java
index 1a75bcfe2feea..cc348ff4ad08a 100644
--- a/presto-mongodb/src/main/java/com/facebook/presto/mongodb/MongoMetadata.java
+++ b/presto-mongodb/src/main/java/com/facebook/presto/mongodb/MongoMetadata.java
@@ -49,7 +49,7 @@
 import java.util.concurrent.atomic.AtomicReference;
 
 import static com.google.common.base.Preconditions.checkState;
-import static java.util.Locale.ENGLISH;
+import static java.util.Locale.ROOT;
 import static java.util.Objects.requireNonNull;
 import static java.util.stream.Collectors.toList;
 
@@ -61,10 +61,12 @@ public class MongoMetadata
     private final MongoSession mongoSession;
 
     private final AtomicReference rollbackAction = new AtomicReference<>();
+    private final MongoClientConfig mongoClientConfig;
 
-    public MongoMetadata(MongoSession mongoSession)
+    public MongoMetadata(MongoSession mongoSession, MongoClientConfig mongoClientConfig)
     {
         this.mongoSession = requireNonNull(mongoSession, "mongoSession is null");
+        this.mongoClientConfig = mongoClientConfig;
     }
 
     @Override
@@ -101,7 +103,7 @@ public List listTables(ConnectorSession session, String schemaN
 
         for (String schemaName : listSchemas(session, schemaNameOrNull)) {
             for (String tableName : mongoSession.getAllTables(schemaName)) {
-                tableNames.add(new SchemaTableName(schemaName, tableName.toLowerCase(ENGLISH)));
+                tableNames.add(new SchemaTableName(schemaName, normalizeIdentifier(session, tableName)));
             }
         }
         return tableNames.build();
@@ -151,7 +153,11 @@ public ColumnMetadata getColumnMetadata(ConnectorSession session, ConnectorTable
     }
 
     @Override
-    public List getTableLayouts(ConnectorSession session, ConnectorTableHandle table, Constraint constraint, Optional> desiredColumns)
+    public ConnectorTableLayoutResult getTableLayoutForConstraint(
+            ConnectorSession session,
+            ConnectorTableHandle table,
+            Constraint constraint,
+            Optional> desiredColumns)
     {
         MongoTableHandle tableHandle = (MongoTableHandle) table;
 
@@ -181,7 +187,7 @@ public List getTableLayouts(ConnectorSession session
                 Optional.empty(),
                 localProperties.build());
 
-        return ImmutableList.of(new ConnectorTableLayoutResult(layout, constraint.getSummary()));
+        return new ConnectorTableLayoutResult(layout, constraint.getSummary());
     }
 
     @Override
@@ -190,8 +196,7 @@ public ConnectorTableLayout getTableLayout(ConnectorSession session, ConnectorTa
         MongoTableLayoutHandle layout = (MongoTableLayoutHandle) handle;
 
         // tables in this connector have a single layout
-        return getTableLayouts(session, layout.getTable(), Constraint.alwaysTrue(), Optional.empty())
-                .get(0)
+        return getTableLayoutForConstraint(session, layout.getTable(), Constraint.alwaysTrue(), Optional.empty())
                 .getTableLayout();
     }
 
@@ -319,4 +324,10 @@ public void dropColumn(ConnectorSession session, ConnectorTableHandle tableHandl
     {
         mongoSession.dropColumn(((MongoTableHandle) tableHandle), ((MongoColumnHandle) column).getName());
     }
+
+    @Override
+    public String normalizeIdentifier(ConnectorSession session, String identifier)
+    {
+        return mongoClientConfig.isCaseSensitiveNameMatchingEnabled() ? identifier : identifier.toLowerCase(ROOT);
+    }
 }
diff --git a/presto-mongodb/src/main/java/com/facebook/presto/mongodb/MongoSession.java b/presto-mongodb/src/main/java/com/facebook/presto/mongodb/MongoSession.java
index c3fb4f1913575..c73f949ac551d 100644
--- a/presto-mongodb/src/main/java/com/facebook/presto/mongodb/MongoSession.java
+++ b/presto-mongodb/src/main/java/com/facebook/presto/mongodb/MongoSession.java
@@ -404,7 +404,7 @@ private Document getTableMetadata(SchemaTableName schemaTableName)
     public boolean collectionExists(MongoDatabase db, String collectionName)
     {
         for (String name : db.listCollectionNames()) {
-            if (name.equalsIgnoreCase(collectionName)) {
+            if (name.equals(collectionName)) {
                 return true;
             }
         }
@@ -457,14 +457,10 @@ private boolean deleteTableMetadata(SchemaTableName schemaTableName)
         String tableName = schemaTableName.getTableName();
 
         MongoDatabase db = client.getDatabase(schemaName);
-        if (!collectionExists(db, tableName)) {
-            return false;
-        }
-
         DeleteResult result = db.getCollection(schemaCollection)
                 .deleteOne(new Document(TABLE_NAME_KEY, tableName));
 
-        return result.getDeletedCount() == 1;
+        return result.getDeletedCount() == 1 || !collectionExists(db, tableName);
     }
 
     private List guessTableFields(SchemaTableName schemaTableName)
diff --git a/presto-mongodb/src/test/java/com/facebook/presto/mongodb/MongoQueryRunner.java b/presto-mongodb/src/test/java/com/facebook/presto/mongodb/MongoQueryRunner.java
index 9c07bed90d15a..1b079c6c93645 100644
--- a/presto-mongodb/src/test/java/com/facebook/presto/mongodb/MongoQueryRunner.java
+++ b/presto-mongodb/src/test/java/com/facebook/presto/mongodb/MongoQueryRunner.java
@@ -24,6 +24,7 @@
 import io.airlift.tpch.TpchTable;
 
 import java.net.InetSocketAddress;
+import java.util.HashMap;
 import java.util.Map;
 
 import static com.facebook.airlift.testing.Closeables.closeAllSuppress;
@@ -53,10 +54,10 @@ private MongoQueryRunner(Session session, int workers)
     public static MongoQueryRunner createMongoQueryRunner(TpchTable... tables)
             throws Exception
     {
-        return createMongoQueryRunner(ImmutableList.copyOf(tables));
+        return createMongoQueryRunner(ImmutableList.copyOf(tables), ImmutableMap.of());
     }
 
-    public static MongoQueryRunner createMongoQueryRunner(Iterable> tables)
+    public static MongoQueryRunner createMongoQueryRunner(Iterable> tables, Map connectorProperties)
             throws Exception
     {
         MongoQueryRunner queryRunner = null;
@@ -66,12 +67,12 @@ public static MongoQueryRunner createMongoQueryRunner(Iterable> tab
             queryRunner.installPlugin(new TpchPlugin());
             queryRunner.createCatalog("tpch", "tpch");
 
-            Map properties = ImmutableMap.of(
-                    "mongodb.seeds", queryRunner.getAddress().getHostString() + ":" + queryRunner.getAddress().getPort(),
-                    "mongodb.socket-keep-alive", "true");
+            connectorProperties = new HashMap<>(connectorProperties);
+            connectorProperties.putIfAbsent("mongodb.seeds", queryRunner.getAddress().getHostString() + ":" + queryRunner.getAddress().getPort());
+            connectorProperties.putIfAbsent("mongodb.socket-keep-alive", "true");
 
             queryRunner.installPlugin(new MongoPlugin());
-            queryRunner.createCatalog("mongodb", "mongodb", properties);
+            queryRunner.createCatalog("mongodb", "mongodb", connectorProperties);
 
             copyTpchTables(queryRunner, "tpch", TINY_SCHEMA_NAME, createSession(), tables);
 
diff --git a/presto-mongodb/src/test/java/com/facebook/presto/mongodb/SyncMemoryBackend.java b/presto-mongodb/src/test/java/com/facebook/presto/mongodb/SyncMemoryBackend.java
index 828230a075424..c8c7abdcc392c 100644
--- a/presto-mongodb/src/test/java/com/facebook/presto/mongodb/SyncMemoryBackend.java
+++ b/presto-mongodb/src/test/java/com/facebook/presto/mongodb/SyncMemoryBackend.java
@@ -13,7 +13,7 @@
  */
 package com.facebook.presto.mongodb;
 
-import de.bwaldvogel.mongo.MongoBackend;
+import de.bwaldvogel.mongo.backend.CursorRegistry;
 import de.bwaldvogel.mongo.backend.memory.MemoryBackend;
 import de.bwaldvogel.mongo.backend.memory.MemoryDatabase;
 import de.bwaldvogel.mongo.exception.MongoServerException;
@@ -25,16 +25,16 @@ public class SyncMemoryBackend
     public MemoryDatabase openOrCreateDatabase(String databaseName)
             throws MongoServerException
     {
-        return new SyncMemoryDatabase(this, databaseName);
+        return new SyncMemoryDatabase(databaseName, this.getCursorRegistry());
     }
 
     private static class SyncMemoryDatabase
             extends MemoryDatabase
     {
-        public SyncMemoryDatabase(MongoBackend backend, String databaseName)
+        public SyncMemoryDatabase(String databaseName, CursorRegistry cursorRegistry)
                 throws MongoServerException
         {
-            super(backend, databaseName);
+            super(databaseName, cursorRegistry);
         }
     }
 }
diff --git a/presto-mongodb/src/test/java/com/facebook/presto/mongodb/TestMongoClientConfig.java b/presto-mongodb/src/test/java/com/facebook/presto/mongodb/TestMongoClientConfig.java
index 42e4c35c2edfe..8c341c1d407d9 100644
--- a/presto-mongodb/src/test/java/com/facebook/presto/mongodb/TestMongoClientConfig.java
+++ b/presto-mongodb/src/test/java/com/facebook/presto/mongodb/TestMongoClientConfig.java
@@ -43,7 +43,8 @@ public void testDefaults()
                 .setReadPreferenceTags("")
                 .setWriteConcern(WriteConcernType.ACKNOWLEDGED)
                 .setRequiredReplicaSetName(null)
-                .setImplicitRowFieldPrefix("_pos"));
+                .setImplicitRowFieldPrefix("_pos")
+                .setCaseSensitiveNameMatchingEnabled(false));
     }
 
     @Test
@@ -66,6 +67,7 @@ public void testExplicitPropertyMappings()
                 .put("mongodb.write-concern", "UNACKNOWLEDGED")
                 .put("mongodb.required-replica-set", "replica_set")
                 .put("mongodb.implicit-row-field-prefix", "_prefix")
+                .put("case-sensitive-name-matching", "true")
                 .build();
 
         MongoClientConfig expected = new MongoClientConfig()
@@ -84,7 +86,8 @@ public void testExplicitPropertyMappings()
                 .setReadPreferenceTags("tag_name:tag_value")
                 .setWriteConcern(WriteConcernType.UNACKNOWLEDGED)
                 .setRequiredReplicaSetName("replica_set")
-                .setImplicitRowFieldPrefix("_prefix");
+                .setImplicitRowFieldPrefix("_prefix")
+                .setCaseSensitiveNameMatchingEnabled(true);
 
         ConfigAssertions.assertFullMapping(properties, expected);
     }
diff --git a/presto-mongodb/src/test/java/com/facebook/presto/mongodb/TestMongoDbIntegrationMixedCaseTest.java b/presto-mongodb/src/test/java/com/facebook/presto/mongodb/TestMongoDbIntegrationMixedCaseTest.java
new file mode 100644
index 0000000000000..24ba8fb81e1bd
--- /dev/null
+++ b/presto-mongodb/src/test/java/com/facebook/presto/mongodb/TestMongoDbIntegrationMixedCaseTest.java
@@ -0,0 +1,223 @@
+/*
+ * Licensed 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 com.facebook.presto.mongodb;
+
+import com.facebook.presto.Session;
+import com.facebook.presto.testing.MaterializedResult;
+import com.facebook.presto.testing.QueryRunner;
+import com.facebook.presto.tests.AbstractTestQueryFramework;
+import com.google.common.collect.ImmutableMap;
+import com.mongodb.MongoClient;
+import io.airlift.tpch.TpchTable;
+import org.bson.Document;
+import org.testng.annotations.AfterClass;
+import org.testng.annotations.BeforeClass;
+import org.testng.annotations.Test;
+
+import java.util.List;
+import java.util.stream.Collectors;
+
+import static com.facebook.presto.common.type.BigintType.BIGINT;
+import static com.facebook.presto.common.type.VarcharType.VARCHAR;
+import static com.facebook.presto.mongodb.MongoQueryRunner.createMongoQueryRunner;
+import static com.facebook.presto.testing.TestingSession.testSessionBuilder;
+import static com.facebook.presto.testing.assertions.Assert.assertEquals;
+import static org.testng.Assert.assertFalse;
+import static org.testng.Assert.assertTrue;
+
+@Test
+public class TestMongoDbIntegrationMixedCaseTest
+        extends AbstractTestQueryFramework
+{
+    private MongoQueryRunner mongoQueryRunner;
+
+    @Override
+    protected QueryRunner createQueryRunner()
+            throws Exception
+    {
+        return createMongoQueryRunner(TpchTable.getTables(), ImmutableMap.of("case-sensitive-name-matching", "true"));
+    }
+
+    @BeforeClass
+    public void setUp()
+    {
+        mongoQueryRunner = (MongoQueryRunner) getQueryRunner();
+    }
+
+    @AfterClass(alwaysRun = true)
+    public final void destroy()
+    {
+        if (mongoQueryRunner != null) {
+            mongoQueryRunner.shutdown();
+        }
+    }
+
+    public void testDescribeTableWithDifferentCaseInSameSchema()
+    {
+        try {
+            getQueryRunner().execute("CREATE TABLE ORDERS AS SELECT * FROM orders");
+
+            assertTrue(getQueryRunner().tableExists(getSession(), "orders"));
+            assertTrue(getQueryRunner().tableExists(getSession(), "ORDERS"));
+
+            MaterializedResult actualColumns = computeActual("DESC ORDERS").toTestTypes();
+
+            MaterializedResult expectedColumns = MaterializedResult.resultBuilder(getQueryRunner().getDefaultSession(), VARCHAR, VARCHAR, VARCHAR, VARCHAR, BIGINT, BIGINT, BIGINT)
+                    .row("orderkey", "bigint", "", "", 19L, null, null)
+                    .row("custkey", "bigint", "", "", 19L, null, null)
+                    .row("orderstatus", "varchar(1)", "", "", null, null, 1L)
+                    .row("totalprice", "double", "", "", 53L, null, null)
+                    .row("orderdate", "date", "", "", null, null, null)
+                    .row("orderpriority", "varchar(15)", "", "", null, null, 15L)
+                    .row("clerk", "varchar(15)", "", "", null, null, 15L)
+                    .row("shippriority", "integer", "", "", 10L, null, null)
+                    .row("comment", "varchar(79)", "", "", null, null, 79L)
+                    .build();
+            assertEquals(actualColumns, expectedColumns);
+        }
+        finally {
+            assertUpdate("DROP TABLE tpch.ORDERS");
+        }
+    }
+
+    @Test
+    public void testCreateAndDropTable()
+    {
+        Session session = testSessionBuilder()
+                .setCatalog("mongodb")
+                .setSchema("Mixed_Test_Database")
+                .build();
+
+        try {
+            getQueryRunner().execute(session, "CREATE TABLE TEST_CREATE(name VARCHAR(50), id int)");
+            assertTrue(getQueryRunner().tableExists(session, "TEST_CREATE"));
+            assertFalse(getQueryRunner().tableExists(session, "test_create"));
+
+            getQueryRunner().execute(session, "CREATE TABLE test_create(name VARCHAR(50), id int)");
+            assertTrue(getQueryRunner().tableExists(session, "test_create"));
+            assertFalse(getQueryRunner().tableExists(session, "Test_Create"));
+        }
+
+        finally {
+            assertUpdate(session, "DROP TABLE TEST_CREATE");
+            assertFalse(getQueryRunner().tableExists(session, "TEST_CREATE"));
+
+            assertUpdate(session, "DROP TABLE test_create");
+            assertFalse(getQueryRunner().tableExists(session, "test_create"));
+        }
+    }
+
+    @Test
+    public void testCreateTableAs()
+    {
+        Session session = testSessionBuilder()
+                .setCatalog("mongodb")
+                .setSchema("Mixed_Test_Database")
+                .build();
+
+        try {
+            getQueryRunner().execute(session, "CREATE TABLE TEST_CTAS AS SELECT * FROM tpch.region");
+            assertTrue(getQueryRunner().tableExists(session, "TEST_CTAS"));
+
+            getQueryRunner().execute(session, "CREATE TABLE IF NOT EXISTS test_ctas AS SELECT * FROM tpch.region");
+            assertTrue(getQueryRunner().tableExists(session, "test_ctas"));
+
+            getQueryRunner().execute(session, "CREATE TABLE TEST_CTAS_Join AS SELECT c.custkey, o.orderkey FROM " +
+                    "tpch.customer c INNER JOIN tpch.orders o ON c.custkey = o.custkey WHERE c.mktsegment = 'BUILDING'");
+            assertTrue(getQueryRunner().tableExists(session, "TEST_CTAS_Join"));
+
+            assertQueryFails("CREATE TABLE Mixed_Test_Database.TEST_CTAS_FAIL_Join AS SELECT c.custkey, o.orderkey FROM " +
+                    "tpch.customer c INNER JOIN tpch.ORDERS1 o ON c.custkey = o.custkey WHERE c.mktsegment = 'BUILDING'", "Table mongodb.tpch.ORDERS1 does not exist"); //failure scenario since tpch.ORDERS1 doesn't exist
+            assertFalse(getQueryRunner().tableExists(session, "TEST_CTAS_FAIL_Join"));
+
+            getQueryRunner().execute(session, "CREATE TABLE Test_CTAS_Mixed_Join AS SELECT Cus.custkey, Ord.orderkey FROM " +
+                    "tpch.customer Cus INNER JOIN tpch.orders Ord ON Cus.custkey = Ord.custkey WHERE Cus.mktsegment = 'BUILDING'");
+            assertTrue(getQueryRunner().tableExists(session, "Test_CTAS_Mixed_Join"));
+        }
+        finally {
+            getQueryRunner().execute(session, "DROP TABLE IF EXISTS TEST_CTAS");
+            getQueryRunner().execute(session, "DROP TABLE IF EXISTS test_ctas");
+            getQueryRunner().execute(session, "DROP TABLE IF EXISTS TEST_CTAS_Join");
+            getQueryRunner().execute(session, "DROP TABLE IF EXISTS Test_CTAS_Mixed_Join");
+        }
+    }
+
+    @Test
+    public void testInsert()
+    {
+        Session session = testSessionBuilder()
+                .setCatalog("mongodb")
+                .setSchema("Mixed_Test_Database")
+                .build();
+
+        try {
+            getQueryRunner().execute(session, "CREATE TABLE Test_Insert (x bigint, y varchar(100))");
+            getQueryRunner().execute(session, "INSERT INTO Test_Insert VALUES (123, 'test')");
+            assertTrue(getQueryRunner().tableExists(session, "Test_Insert"));
+            assertQuery("SELECT * FROM Mixed_Test_Database.Test_Insert", "SELECT 123 x, 'test' y");
+
+            getQueryRunner().execute(session, "CREATE TABLE TEST_INSERT (x bigint, Y varchar(100))");
+            getQueryRunner().execute(session, "INSERT INTO TEST_INSERT VALUES (1234, 'test1')");
+            assertTrue(getQueryRunner().tableExists(session, "TEST_INSERT"));
+            assertQuery("SELECT * FROM Mixed_Test_Database.TEST_INSERT", "SELECT 1234 x, 'test1' Y");
+        }
+        finally {
+            getQueryRunner().execute(session, "DROP TABLE IF EXISTS Test_Insert");
+            getQueryRunner().execute(session, "DROP TABLE IF EXISTS TEST_INSERT");
+        }
+    }
+
+    @Test
+    public void testMixedCaseColumns()
+    {
+        try {
+            assertUpdate("CREATE TABLE test (a integer, B integer)");
+            assertUpdate("CREATE TABLE TEST (a integer, aA integer)");
+            assertUpdate("INSERT INTO TEST VALUES (123, 12)", 1);
+            assertTableColumnNames("TEST", "a", "aA");
+            assertUpdate("ALTER TABLE TEST ADD COLUMN EMAIL varchar");
+            assertUpdate("ALTER TABLE TEST RENAME COLUMN a TO a_New");
+            assertTableColumnNames("TEST", "a_New", "aA", "EMAIL");
+            assertUpdate("ALTER TABLE TEST DROP COLUMN aA");
+            assertTableColumnNames("TEST", "a_New", "EMAIL");
+        }
+        finally {
+            assertUpdate("DROP TABLE test");
+            assertUpdate("DROP TABLE TEST");
+        }
+    }
+
+    @Test
+    public void testShowSchemas()
+    {
+        // Create two MongoDB databases directly, since Presto doesn't support create schema for mongodb
+        MongoClient mongoClient = mongoQueryRunner.getMongoClient();
+        try {
+            mongoClient.getDatabase("TESTDB1").getCollection("dummy").insertOne(new Document("x", 1));
+            mongoClient.getDatabase("testdb1").getCollection("dummy").insertOne(new Document("x", 1));
+
+            MaterializedResult result = computeActual("SHOW SCHEMAS");
+            List schemaNames = result.getMaterializedRows().stream()
+                    .map(row -> row.getField(0).toString())
+                    .collect(Collectors.toList());
+
+            assertTrue(schemaNames.contains("TESTDB1"));
+            assertTrue(schemaNames.contains("testdb1"));
+        }
+        finally {
+            mongoClient.getDatabase("TESTDB1").drop();
+            mongoClient.getDatabase("testdb1").drop();
+        }
+    }
+}
diff --git a/presto-mongodb/src/test/java/com/facebook/presto/mongodb/TestMongoDistributedQueries.java b/presto-mongodb/src/test/java/com/facebook/presto/mongodb/TestMongoDistributedQueries.java
index 08e6e9f04feac..23c58a846d886 100644
--- a/presto-mongodb/src/test/java/com/facebook/presto/mongodb/TestMongoDistributedQueries.java
+++ b/presto-mongodb/src/test/java/com/facebook/presto/mongodb/TestMongoDistributedQueries.java
@@ -15,6 +15,7 @@
 
 import com.facebook.presto.testing.QueryRunner;
 import com.facebook.presto.tests.AbstractTestQueries;
+import com.google.common.collect.ImmutableMap;
 import io.airlift.tpch.TpchTable;
 import org.testng.annotations.AfterClass;
 import org.testng.annotations.BeforeClass;
@@ -32,7 +33,7 @@ public class TestMongoDistributedQueries
     protected QueryRunner createQueryRunner()
             throws Exception
     {
-        return createMongoQueryRunner(TpchTable.getTables());
+        return createMongoQueryRunner(TpchTable.getTables(), ImmutableMap.of());
     }
 
     @BeforeClass
diff --git a/presto-mongodb/src/test/java/com/facebook/presto/mongodb/TestMongoIntegrationSmokeTest.java b/presto-mongodb/src/test/java/com/facebook/presto/mongodb/TestMongoIntegrationSmokeTest.java
index a5e820d4a1dfb..fd35e5ea284ee 100644
--- a/presto-mongodb/src/test/java/com/facebook/presto/mongodb/TestMongoIntegrationSmokeTest.java
+++ b/presto-mongodb/src/test/java/com/facebook/presto/mongodb/TestMongoIntegrationSmokeTest.java
@@ -21,6 +21,7 @@
 import com.google.common.collect.ImmutableMap;
 import com.mongodb.client.MongoCollection;
 import org.bson.Document;
+import org.bson.types.ObjectId;
 import org.testng.annotations.AfterClass;
 import org.testng.annotations.BeforeClass;
 import org.testng.annotations.Test;
@@ -146,7 +147,8 @@ public void testInsertWithEveryType()
         assertEquals(row.getField(5), LocalDate.of(1980, 5, 7));
         assertEquals(row.getField(6), LocalDateTime.of(1980, 5, 7, 11, 22, 33, 456_000_000));
         assertEquals(row.getField(7), LocalTime.of(11, 22, 33, 456_000_000));
-        assertEquals(row.getField(8), "{\"name\":\"alice\"}");
+        assertEquals(new ObjectId((byte[]) row.getField(8)), new ObjectId("ffffffffffffffffffffffff"));
+        assertEquals(row.getField(9).toString(), "{\"name\":\"alice\"}");
         assertUpdate("DROP TABLE test_insert_types_table");
         assertFalse(getQueryRunner().tableExists(getSession(), "test_insert_types_table"));
     }
@@ -218,13 +220,13 @@ public void testMaps()
         assertUpdate("CREATE TABLE test.tmp_map9 (col VARCHAR)");
         mongoQueryRunner.getMongoClient().getDatabase("test").getCollection("tmp_map9").insertOne(new Document(
                 ImmutableMap.of("col", new Document(ImmutableMap.of("key1", "value1", "key2", "value2")))));
-        assertQuery("SELECT col FROM test.tmp_map9", "SELECT '{ \"key1\" : \"value1\", \"key2\" : \"value2\" }'");
+        assertQuery("SELECT col FROM test.tmp_map9", "SELECT '{\"key1\": \"value1\", \"key2\": \"value2\"}'");
 
         assertUpdate("CREATE TABLE test.tmp_map10 (col VARCHAR)");
         mongoQueryRunner.getMongoClient().getDatabase("test").getCollection("tmp_map10").insertOne(new Document(
                 ImmutableMap.of("col", ImmutableList.of(new Document(ImmutableMap.of("key1", "value1", "key2", "value2")),
                         new Document(ImmutableMap.of("key3", "value3", "key4", "value4"))))));
-        assertQuery("SELECT col FROM test.tmp_map10", "SELECT '[{ \"key1\" : \"value1\", \"key2\" : \"value2\" }, { \"key3\" : \"value3\", \"key4\" : \"value4\" }]'");
+        assertQuery("SELECT col FROM test.tmp_map10", "SELECT '[{\"key1\": \"value1\", \"key2\": \"value2\"}, {\"key3\": \"value3\", \"key4\": \"value4\"}]'");
 
         assertUpdate("CREATE TABLE test.tmp_map11 (col VARCHAR)");
         mongoQueryRunner.getMongoClient().getDatabase("test").getCollection("tmp_map11").insertOne(new Document(
diff --git a/presto-mysql/src/test/java/com/facebook/presto/plugin/mysql/TestMySqlIntegrationSmokeTest.java b/presto-mysql/src/test/java/com/facebook/presto/plugin/mysql/TestMySqlIntegrationSmokeTest.java
index adbbd0b928a35..65b97f1788d7d 100644
--- a/presto-mysql/src/test/java/com/facebook/presto/plugin/mysql/TestMySqlIntegrationSmokeTest.java
+++ b/presto-mysql/src/test/java/com/facebook/presto/plugin/mysql/TestMySqlIntegrationSmokeTest.java
@@ -216,6 +216,46 @@ public void testMysqlGeometry()
         assertUpdate("DROP TABLE tpch.test_geometry");
     }
 
+    @Test
+    public void testMysqlDecimal()
+    {
+        assertUpdate("CREATE TABLE test_decimal (d DECIMAL(10, 2))");
+
+        assertUpdate("INSERT INTO test_decimal VALUES (123.45)", 1);
+        assertUpdate("INSERT INTO test_decimal VALUES (67890.12)", 1);
+        assertUpdate("INSERT INTO test_decimal VALUES (0.99)", 1);
+
+        assertQuery(
+                "SELECT d FROM test_decimal WHERE d<200.00 AND d>0.00",
+                "VALUES " +
+                        "CAST('123.45' AS DECIMAL), " +
+                        "CAST('0.99' AS DECIMAL)");
+
+        assertUpdate("DROP TABLE test_decimal");
+    }
+
+    @Test
+    public void testMysqlTime()
+    {
+        assertUpdate("CREATE TABLE test_time (datatype_time time)");
+
+        assertUpdate("INSERT INTO test_time VALUES (time '01:02:03.456')", 1);
+
+        assertQuery(
+                "SELECT datatype_time FROM test_time",
+                "VALUES " +
+                        "CAST('01:02:03.456' AS time)");
+
+        assertUpdate("DROP TABLE test_time");
+    }
+
+    @Test
+    public void testMysqlUnsupportedTimeTypes()
+    {
+        assertQueryFails("CREATE TABLE test_timestamp_with_timezone (timestamp_with_time_zone timestamp with time zone)", "Unsupported column type: timestamp with time zone");
+        assertQueryFails("CREATE TABLE test_time_with_timezone (time_with_with_time_zone time with time zone)", "Unsupported column type: time with time zone");
+    }
+
     @Test
     public void testCharTrailingSpace()
             throws Exception
diff --git a/presto-native-execution/CMakeLists.txt b/presto-native-execution/CMakeLists.txt
index b85d27d8c691e..1e7fce6c7717b 100644
--- a/presto-native-execution/CMakeLists.txt
+++ b/presto-native-execution/CMakeLists.txt
@@ -33,6 +33,12 @@ set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} ${SCRIPT_CXX_FLAGS}")
 set(DISABLED_WARNINGS
     "-Wno-nullability-completeness -Wno-deprecated-declarations")
 
+if(CMAKE_CXX_COMPILER_ID STREQUAL "GNU")
+  if(CMAKE_CXX_COMPILER_VERSION VERSION_GREATER_EQUAL "14.0.0")
+    string(APPEND DISABLED_WARNINGS " -Wno-error=template-id-cdtor")
+  endif()
+endif()
+
 # Important warnings that must be explicitly enabled.
 set(ENABLE_WARNINGS "-Wreorder")
 
@@ -77,58 +83,46 @@ option(PRESTO_ENABLE_SPATIAL "Enable spatial support" ON)
 # and turn on int128.
 add_compile_definitions(FOLLY_HAVE_INT128_T=1 FOLLY_CFG_NO_COROUTINES)
 
-if(PRESTO_ENABLE_S3)
-  set(VELOX_ENABLE_S3
-      ON
-      CACHE BOOL "Build S3 support")
-endif()
+set(VELOX_ENABLE_S3
+    ${PRESTO_ENABLE_S3}
+    CACHE BOOL "Build S3 support")
 
-if(PRESTO_ENABLE_HDFS)
-  set(VELOX_ENABLE_HDFS
-      ON
-      CACHE BOOL "Build HDFS support")
-endif()
+set(VELOX_ENABLE_HDFS
+    ${PRESTO_ENABLE_HDFS}
+    CACHE BOOL "Build HDFS support")
 
-if(PRESTO_ENABLE_GCS)
-  set(VELOX_ENABLE_GCS
-      ON
-      CACHE BOOL "Build GCS support")
-endif()
+set(VELOX_ENABLE_GCS
+    ${PRESTO_ENABLE_GCS}
+    CACHE BOOL "Build GCS support")
 
-if(PRESTO_ENABLE_ABFS)
-  set(VELOX_ENABLE_ABFS
-      ON
-      CACHE BOOL "Build ABFS support")
-endif()
+set(VELOX_ENABLE_ABFS
+    ${PRESTO_ENABLE_ABFS}
+    CACHE BOOL "Build ABFS support")
 
-if(PRESTO_ENABLE_PARQUET)
-  set(VELOX_ENABLE_PARQUET
-      ON
-      CACHE BOOL "Enable Parquet support")
-endif()
+set(VELOX_ENABLE_PARQUET
+    ${PRESTO_ENABLE_PARQUET}
+    CACHE BOOL "Enable Parquet support")
 
+set(VELOX_ENABLE_REMOTE_FUNCTIONS
+    ${PRESTO_ENABLE_REMOTE_FUNCTIONS}
+    CACHE BOOL "Enable remote function support in Velox")
 if(PRESTO_ENABLE_REMOTE_FUNCTIONS)
-  set(VELOX_ENABLE_REMOTE_FUNCTIONS
-      ON
-      CACHE BOOL "Enable remote function support in Velox")
   add_compile_definitions(PRESTO_ENABLE_REMOTE_FUNCTIONS)
 endif()
 
+set(VELOX_ENABLE_CUDF
+    ${PRESTO_ENABLE_CUDF}
+    CACHE BOOL "Enable cuDF support")
 if(PRESTO_ENABLE_CUDF)
-  set(VELOX_ENABLE_CUDF
-      ON
-      CACHE BOOL "Enable cuDF support")
   add_compile_definitions(PRESTO_ENABLE_CUDF)
   enable_language(CUDA)
   # Determine CUDA_ARCHITECTURES automatically.
   cmake_policy(SET CMP0104 NEW)
 endif()
 
-if(PRESTO_ENABLE_SPATIAL)
-  set(VELOX_ENABLE_GEO
-      ON
-      CACHE BOOL "Enable Velox Geometry (aka spatial) support")
-endif()
+set(VELOX_ENABLE_GEO
+    ${PRESTO_ENABLE_SPATIAL}
+    CACHE BOOL "Enable Velox Geometry (aka spatial) support")
 
 set(VELOX_BUILD_TESTING
     OFF
@@ -221,7 +215,7 @@ include_directories(SYSTEM ${FBTHRIFT_INCLUDE_DIR})
 set(PROXYGEN_LIBRARIES ${PROXYGEN_HTTP_SERVER} ${PROXYGEN} ${WANGLE} ${FIZZ}
                        ${MVFST_EXCEPTION})
 find_path(PROXYGEN_DIR NAMES include/proxygen)
-set(PROXYGEN_INCLUDE_DIR "${PROXYGEN_DIR}/include/proxygen")
+set(PROXYGEN_INCLUDE_DIR "${PROXYGEN_DIR}/include/")
 
 include_directories(SYSTEM ${OPENSSL_INCLUDE_DIR} ${PROXYGEN_INCLUDE_DIR})
 include_directories(.)
@@ -233,6 +227,11 @@ include_directories(${CMAKE_BINARY_DIR})
 # set this for backwards compatibility, will be overwritten in velox/
 set(VELOX_GTEST_INCUDE_DIR "velox/third_party/googletest/googletest/include")
 
+# Do not use the Mono library because it causes link errors.
+set(VELOX_MONO_LIBRARY
+    OFF
+    CACHE BOOL "Build Velox mono library")
+
 add_subdirectory(velox)
 
 if(PRESTO_ENABLE_TESTING)
diff --git a/presto-native-execution/README.md b/presto-native-execution/README.md
index f1e31632de56e..5617ab2510676 100644
--- a/presto-native-execution/README.md
+++ b/presto-native-execution/README.md
@@ -95,7 +95,9 @@ S3 support needs the [AWS SDK C++](https://github.com/aws/aws-sdk-cpp) library.
 This dependency can be installed by running the target platform build script
 from the `presto/presto-native-execution` directory.
 
-`./velox/scripts/setup-centos9.sh install_aws`
+`./velox/scripts/setup-centos9.sh install_aws_deps`
+    Or
+`./velox/scripts/setup-ubuntu.sh install_aws_deps`
 
 #### JWT Authentication
 To enable JWT authentication support, add `-DPRESTO_ENABLE_JWT=ON` to the
@@ -214,12 +216,23 @@ Run IcebergExternalWorkerQueryRunner,
       `$DATA_DIR/iceberg_data//`. Here `file_format` could be `PARQUET | ORC | AVRO` and `catalog_type` could be `HIVE | HADOOP | NESSIE | REST`.
   * Use classpath of module: choose `presto-native-execution` module.
 
+Run NativeSidecarPluginQueryRunner:
+* Edit/Create `NativeSidecarPluginQueryRunner` Application Run/Debug Configuration (alter paths accordingly).
+  * Main class: `com.facebook.presto.sidecar.NativeSidecarPluginQueryRunner`.
+  * VM options : `-ea -Xmx5G -XX:+ExitOnOutOfMemoryError -Duser.timezone=America/Bahia_Banderas -Dhive.security=legacy`.
+  * Working directory: `$MODULE_DIR$`
+  * Environment variables: `PRESTO_SERVER=/Users//git/presto/presto-native-execution/cmake-build-debug/presto_cpp/main/presto_server;DATA_DIR=/Users//Desktop/data;WORKER_COUNT=0`
+  * Use classpath of module: choose `presto-native-sidecar-plugin` module.
+
 Run CLion:
 * File->Close Project if any is open.
 * Open `presto/presto-native-execution` directory as CMake project and wait till CLion loads/generates cmake files, symbols, etc.
 * Edit configuration for `presto_server` module (alter paths accordingly).
   * Program arguments: `--logtostderr=1 --v=1 --etc_dir=/Users//git/presto/presto-native-execution/etc`
   * Working directory: `/Users//git/presto/presto-native-execution`
+* For sidecar, Edit configuration for `presto_server` module (alter paths accordingly).
+  * Program arguments: `--logtostderr=1 --v=1 --etc_dir=/Users//git/presto/presto-native-execution/etc_sidecar`
+  * Working directory: `/Users//git/presto/presto-native-execution`
 * Edit menu CLion->Preferences->Build, Execution, Deployment->CMake
   * CMake options: `-DVELOX_BUILD_TESTING=ON -DCMAKE_BUILD_TYPE=Debug`
   * Build options: `-- -j 12`
@@ -229,7 +242,7 @@ Run CLion:
 * To enable clang format you need
   * Open any h or cpp file in the editor and select `Enable ClangFormat` by clicking `4 spaces` rectangle in the status bar (bottom right) which is next to `UTF-8` bar.
 
-    ![ScreenShot](cl_clangformat_switcherenable.png)
+    ![ScreenShot](docs/images/cl_clangformat_switcherenable.png)
 
 ### Run Presto Coordinator + Worker
 * Note that everything below can be done without using IDEs by running command line commands (not in this readme).
@@ -237,12 +250,24 @@ Run CLion:
   * For Hive, Run `HiveExternalWorkerQueryRunner` from IntelliJ and wait until it starts (`======== SERVER STARTED ========` is displayed in the log output).
   * For Iceberg, Run `IcebergExternalWorkerQueryRunner` from IntelliJ and wait until it starts (`======== SERVER STARTED ========` is displayed in the log output).
 * Scroll up the log output and find `Discovery URL http://127.0.0.1:50555`. The port is 'random' with every start.
-* Copy that port (or the whole URL) to the `discovery.uri` field in `presto/presto-native-execution/etc/config.properties` for the worker to discover the Coordinator.
+* Copy that port (or the whole URL) to the `discovery.uri` field in `presto/presto-native-execution/etc/config.properties` for the worker to announce itself to the Coordinator.
 * In CLion run "presto_server" module. Connection success will be indicated by `Announcement succeeded: 202` line in the log output.
-* Two ways to run Presto client to start executing queries on the running local setup:
-  1. In command line from presto root directory run the presto client:
-      * `java -jar presto-cli/target/presto-cli-*-executable.jar --catalog hive --schema tpch`
-  2. Run `Presto Client` Application (see above on how to create and setup the configuration) inside IntelliJ
+* See **Run Presto Client** to start executing queries on the running local setup.
+
+### Run Presto Coordinator + Sidecar
+* Note that everything below can be done without using IDEs by running command line commands (not in this readme).
+* Add a property `presto.default-namespace=native.default` to `presto-native-execution/etc/config.properties`.
+* Run `NativeSidecarPluginQueryRunner` from IntelliJ and wait until it starts (`======== SERVER STARTED ========` is displayed in the log output).
+* Scroll up the log output and find `Discovery URL http://127.0.0.1:50555`. The port is 'random' with every startup.
+* Copy that port (or the whole URL) to the `discovery.uri` field in`presto/presto-native-execution/etc_sidecar/config.properties` for the sidecar to announce itself to the Coordinator.
+* In CLion run "presto_server" module. Connection success will be indicated by `Announcement succeeded: 202` line in the log output.
+* See **Run Presto Client** to start executing queries on the running local setup.
+
+### Run Presto Client
+* Run the following command from the presto root directory to start the Presto client:
+    ```
+    java -jar presto-cli/target/presto-cli-*-executable.jar --catalog hive --schema tpch
+    ```
 * You can start from `show tables;` and `describe table;` queries and execute more queries as needed.
 
 ### Run Integration (End to End or E2E) Tests
diff --git a/presto-native-execution/cl_clangformat_switcherenable.png b/presto-native-execution/docs/images/cl_clangformat_switcherenable.png
similarity index 100%
rename from presto-native-execution/cl_clangformat_switcherenable.png
rename to presto-native-execution/docs/images/cl_clangformat_switcherenable.png
diff --git a/presto-native-execution/etc/catalog/tpchstandard.properties b/presto-native-execution/etc/catalog/tpchstandard.properties
index 16e833ca8f436..75110c5acf145 100644
--- a/presto-native-execution/etc/catalog/tpchstandard.properties
+++ b/presto-native-execution/etc/catalog/tpchstandard.properties
@@ -1 +1 @@
-connector.name=tpch
\ No newline at end of file
+connector.name=tpch
diff --git a/presto-native-execution/etc/config.properties b/presto-native-execution/etc/config.properties
index 37db096a9402a..bec551886110b 100644
--- a/presto-native-execution/etc/config.properties
+++ b/presto-native-execution/etc/config.properties
@@ -1,4 +1,4 @@
-discovery.uri=http://127.0.0.1:58215
+discovery.uri=http://127.0.0.1:
 presto.version=testversion
 http-server.http.port=7777
 shutdown-onset-sec=1
diff --git a/presto-native-execution/etc/velox.properties b/presto-native-execution/etc/velox.properties
deleted file mode 100644
index 6c2506bd99a8e..0000000000000
--- a/presto-native-execution/etc/velox.properties
+++ /dev/null
@@ -1 +0,0 @@
-mutable-config=true
\ No newline at end of file
diff --git a/presto-native-execution/etc_sidecar/catalog/hive.properties b/presto-native-execution/etc_sidecar/catalog/hive.properties
new file mode 100644
index 0000000000000..466b7e664e44f
--- /dev/null
+++ b/presto-native-execution/etc_sidecar/catalog/hive.properties
@@ -0,0 +1 @@
+connector.name=hive
diff --git a/presto-native-execution/etc_sidecar/catalog/iceberg.properties b/presto-native-execution/etc_sidecar/catalog/iceberg.properties
new file mode 100644
index 0000000000000..f3a43dcb28126
--- /dev/null
+++ b/presto-native-execution/etc_sidecar/catalog/iceberg.properties
@@ -0,0 +1 @@
+connector.name=iceberg
diff --git a/presto-native-execution/etc_sidecar/catalog/tpchstandard.properties b/presto-native-execution/etc_sidecar/catalog/tpchstandard.properties
new file mode 100644
index 0000000000000..75110c5acf145
--- /dev/null
+++ b/presto-native-execution/etc_sidecar/catalog/tpchstandard.properties
@@ -0,0 +1 @@
+connector.name=tpch
diff --git a/presto-native-execution/etc_sidecar/config.properties b/presto-native-execution/etc_sidecar/config.properties
new file mode 100644
index 0000000000000..76ea4efa7c059
--- /dev/null
+++ b/presto-native-execution/etc_sidecar/config.properties
@@ -0,0 +1,7 @@
+discovery.uri=http://127.0.0.1:
+presto.version=testversion
+http-server.http.port=7778
+shutdown-onset-sec=1
+runtime-metrics-collection-enabled=true
+native-sidecar=true
+presto.default-namespace=native.default
diff --git a/presto-native-execution/etc_sidecar/node.properties b/presto-native-execution/etc_sidecar/node.properties
new file mode 100644
index 0000000000000..1d92b7ace8087
--- /dev/null
+++ b/presto-native-execution/etc_sidecar/node.properties
@@ -0,0 +1,3 @@
+node.environment=testing
+node.internal-address=127.0.0.1
+node.location=testing-location
diff --git a/presto-native-execution/pom.xml b/presto-native-execution/pom.xml
index ab6d81aba382b..f9150f63a4daf 100644
--- a/presto-native-execution/pom.xml
+++ b/presto-native-execution/pom.xml
@@ -30,12 +30,37 @@
     
 
     
+        
+        
+            javax.ws.rs
+            javax.ws.rs-api
+            test
+        
+        
+            javax.servlet
+            javax.servlet-api
+            3.1.0
+            test
+        
+        
         
             com.facebook.presto
             presto-tpch
             test
         
 
+        
+            org.jetbrains
+            annotations
+            test
+        
+
+        
+            org.weakref
+            jmxutils
+            test
+        
+
         
             io.airlift.tpch
             tpch
@@ -116,6 +141,22 @@
             test
         
 
+        
+            com.facebook.presto
+            presto-parser
+            test
+            
+                
+                    com.google.guava
+                    guava
+                
+                
+                    com.facebook.presto
+                    presto-spi
+                
+            
+        
+
         
             com.facebook.presto
             presto-iceberg
@@ -300,6 +341,19 @@
             presto-jdbc
             test
         
+
+        
+            org.apache.commons
+            commons-lang3
+            test
+        
+
+        
+            com.facebook.presto
+            presto-sql-invoked-functions-plugin
+            ${project.version}
+            test
+        
     
 
     
diff --git a/presto-native-execution/presto_cpp/main/CMakeLists.txt b/presto-native-execution/presto_cpp/main/CMakeLists.txt
index 38770def39269..29836ba760385 100644
--- a/presto-native-execution/presto_cpp/main/CMakeLists.txt
+++ b/presto-native-execution/presto_cpp/main/CMakeLists.txt
@@ -78,10 +78,10 @@ target_link_libraries(
   velox_hive_iceberg_splitreader
   velox_hive_partition_function
   velox_presto_serializer
+  velox_presto_type_parser
   velox_s3fs
   velox_serialization
   velox_time
-  velox_type_parser
   velox_type
   velox_type_fbhive
   velox_type_tz
diff --git a/presto-native-execution/presto_cpp/main/PrestoServer.cpp b/presto-native-execution/presto_cpp/main/PrestoServer.cpp
index c89f91de77957..75a2bb1eedfe4 100644
--- a/presto-native-execution/presto_cpp/main/PrestoServer.cpp
+++ b/presto-native-execution/presto_cpp/main/PrestoServer.cpp
@@ -354,6 +354,14 @@ void PrestoServer::run() {
         json infoStateJson = convertNodeState(server->nodeState());
         http::sendOkResponse(downstream, infoStateJson);
       });
+  httpServer_->registerGet(
+      "/v1/info/stats",
+      [server = this](
+          proxygen::HTTPMessage* /*message*/,
+          const std::vector>& /*body*/,
+          proxygen::ResponseHandler* downstream) {
+        server->reportNodeStats(downstream);
+      });
   httpServer_->registerPut(
       "/v1/info/state",
       [server = this](
@@ -1743,4 +1751,16 @@ void PrestoServer::createTaskManager() {
       driverExecutor_.get(), httpSrvCpuExecutor_.get(), spillerExecutor_.get());
 }
 
+void PrestoServer::reportNodeStats(proxygen::ResponseHandler* downstream) {
+  protocol::NodeStats nodeStats;
+
+  auto loadMetrics = std::make_shared();
+  loadMetrics->cpuOverload = cpuOverloaded_;
+  loadMetrics->memoryOverload = memOverloaded_;
+
+  nodeStats.loadMetrics = loadMetrics;
+  nodeStats.nodeState = convertNodeState(this->nodeState());
+
+  http::sendOkResponse(downstream, json(nodeStats));
+}
 } // namespace facebook::presto
diff --git a/presto-native-execution/presto_cpp/main/PrestoServer.h b/presto-native-execution/presto_cpp/main/PrestoServer.h
index 00c7697218f4e..8e658c89a157e 100644
--- a/presto-native-execution/presto_cpp/main/PrestoServer.h
+++ b/presto-native-execution/presto_cpp/main/PrestoServer.h
@@ -172,8 +172,8 @@ class PrestoServer {
 
   /// Invoked to get the ip address of the process. In certain deployment
   /// setup, each process has different ip address. Deployment environment
-  /// may provide there own library to get process specific ip address.
-  /// In such cases, getLocalIp can be overriden to pass process specific
+  /// may provide their own library to get process specific ip address.
+  /// In such cases, getLocalIp can be overridden to pass process specific
   /// ip address.
   virtual std::string getLocalIp() const;
 
@@ -208,6 +208,8 @@ class PrestoServer {
 
   void reportNodeStatus(proxygen::ResponseHandler* downstream);
 
+  void reportNodeStats(proxygen::ResponseHandler* downstream);
+
   void handleGracefulShutdown(
       const std::vector>& body,
       proxygen::ResponseHandler* downstream);
diff --git a/presto-native-execution/presto_cpp/main/PrestoTask.cpp b/presto-native-execution/presto_cpp/main/PrestoTask.cpp
index 3b0cb9f03006f..3621a763bdbdb 100644
--- a/presto-native-execution/presto_cpp/main/PrestoTask.cpp
+++ b/presto-native-execution/presto_cpp/main/PrestoTask.cpp
@@ -769,6 +769,8 @@ void PrestoTask::updateTimeInfoLocked(
     taskRuntimeStats["endTime"].addValue(veloxTaskStats.endTimeMs);
   }
   taskRuntimeStats.insert({"nativeProcessCpuTime", fromNanos(processCpuTime_)});
+  // Represents the time between receiving first taskUpdate and task creation time
+  taskRuntimeStats.insert({"taskCreationTime", fromNanos((createFinishTimeMs - firstTimeReceiveTaskUpdateMs) * 1'000'000)});
 }
 
 void PrestoTask::updateMemoryInfoLocked(
diff --git a/presto-native-execution/presto_cpp/main/PrestoTask.h b/presto-native-execution/presto_cpp/main/PrestoTask.h
index 297f877f3793e..b41666fa90a93 100644
--- a/presto-native-execution/presto_cpp/main/PrestoTask.h
+++ b/presto-native-execution/presto_cpp/main/PrestoTask.h
@@ -121,7 +121,12 @@ struct PrestoTask {
   uint64_t lastTaskStatsUpdateMs{0};
 
   uint64_t lastMemoryReservation{0};
+  /// Time point (in ms) when the time we start task creating.
   uint64_t createTimeMs{0};
+  /// Time point (in ms) when the first time we receive task update.
+  uint64_t firstTimeReceiveTaskUpdateMs{0};
+  /// Time point (in ms) when the time we finish task creating.
+  uint64_t createFinishTimeMs{0};
   uint64_t startTimeMs{0};
   uint64_t firstSplitStartTimeMs{0};
   uint64_t lastEndTimeMs{0};
diff --git a/presto-native-execution/presto_cpp/main/PrestoToVeloxQueryConfig.cpp b/presto-native-execution/presto_cpp/main/PrestoToVeloxQueryConfig.cpp
index 8f2fbb5d202ec..e4e216450287a 100644
--- a/presto-native-execution/presto_cpp/main/PrestoToVeloxQueryConfig.cpp
+++ b/presto-native-execution/presto_cpp/main/PrestoToVeloxQueryConfig.cpp
@@ -139,6 +139,12 @@ void updateFromSystemConfigs(
       {std::string(SystemConfig::kUseLegacyArrayAgg),
        velox::core::QueryConfig::kPrestoArrayAggIgnoreNulls},
 
+      {std::string{SystemConfig::kTaskWriterCount},
+        velox::core::QueryConfig::kTaskWriterCount},
+
+      {std::string{SystemConfig::kTaskPartitionedWriterCount},
+        velox::core::QueryConfig::kTaskPartitionedWriterCount},
+
       {std::string(SystemConfig::kSinkMaxBufferSize),
        velox::core::QueryConfig::kMaxOutputBufferSize,
        [](const auto& value) {
diff --git a/presto-native-execution/presto_cpp/main/QueryContextManager.cpp b/presto-native-execution/presto_cpp/main/QueryContextManager.cpp
index 6387be0fd87fe..b96f03dce2909 100644
--- a/presto-native-execution/presto_cpp/main/QueryContextManager.cpp
+++ b/presto-native-execution/presto_cpp/main/QueryContextManager.cpp
@@ -63,7 +63,8 @@ std::shared_ptr QueryContextManager::createAndCacheQueryCtx(
     QueryContextCache& cache,
     const QueryId& queryId,
     velox::core::QueryConfig&& queryConfig,
-    std::unordered_map>&& connectorConfigs,
+    std::unordered_map>&&
+        connectorConfigs,
     std::shared_ptr&& pool) {
   auto queryCtx = core::QueryCtx::create(
       driverExecutor_,
@@ -96,10 +97,12 @@ std::shared_ptr QueryContextManager::findOrCreateQueryCtx(
   // is still indexed by the query id.
   static std::atomic_uint64_t poolId{0};
   std::optional poolDbgOpts;
-  const auto debugMemoryPoolNameRegex = queryConfig.debugMemoryPoolNameRegex();
+  auto debugMemoryPoolNameRegex = queryConfig.debugMemoryPoolNameRegex();
   if (!debugMemoryPoolNameRegex.empty()) {
     poolDbgOpts = memory::MemoryPool::DebugOptions{
-        .debugPoolNameRegex = debugMemoryPoolNameRegex};
+        .debugPoolNameRegex = std::move(debugMemoryPoolNameRegex),
+        .debugPoolWarnThresholdBytes =
+            queryConfig.debugMemoryPoolWarnThresholdBytes()};
   }
   auto pool = memory::MemoryManager::getInstance()->addRootPool(
       fmt::format("{}_{}", queryId, poolId++),
diff --git a/presto-native-execution/presto_cpp/main/SessionProperties.cpp b/presto-native-execution/presto_cpp/main/SessionProperties.cpp
index e3a03a152b121..d87ff9167cafc 100644
--- a/presto-native-execution/presto_cpp/main/SessionProperties.cpp
+++ b/presto-native-execution/presto_cpp/main/SessionProperties.cpp
@@ -280,6 +280,19 @@ SessionProperties::SessionProperties() {
       QueryConfig::kDebugMemoryPoolNameRegex,
       c.debugMemoryPoolNameRegex());
 
+  addSessionProperty(
+      kDebugMemoryPoolWarnThresholdBytes,
+      "Warning threshold in bytes for debug memory pools. When set to a "
+      "non-zero value, a warning will be logged once per memory pool when "
+      "allocations cause the pool to exceed this threshold. This is useful for "
+      "identifying memory usage patterns during debugging. Requires allocation "
+      "tracking to be enabled with `native_debug_memory_pool_name_regex` "
+      "for the pool. A value of 0 means no warning threshold is enforced.",
+      BIGINT(),
+      false,
+      QueryConfig::kDebugMemoryPoolWarnThresholdBytes,
+      std::to_string(c.debugMemoryPoolWarnThresholdBytes()));
+
   addSessionProperty(
       kSelectiveNimbleReaderEnabled,
       "Temporary flag to control whether selective Nimble reader should be "
diff --git a/presto-native-execution/presto_cpp/main/SessionProperties.h b/presto-native-execution/presto_cpp/main/SessionProperties.h
index 10758ee4c95d8..ab888852d5e00 100644
--- a/presto-native-execution/presto_cpp/main/SessionProperties.h
+++ b/presto-native-execution/presto_cpp/main/SessionProperties.h
@@ -184,6 +184,12 @@ class SessionProperties {
   static constexpr const char* kDebugMemoryPoolNameRegex =
       "native_debug_memory_pool_name_regex";
 
+  /// Warning threshold in bytes for memory pool allocations. Logs callsites 
+  /// when exceeded. Requires allocation tracking to be enabled with 
+  /// `native_debug_memory_pool_name_regex` property for the pool.
+  static constexpr const char* kDebugMemoryPoolWarnThresholdBytes =
+      "native_debug_memory_pool_warn_threshold_bytes";
+
   /// Temporary flag to control whether selective Nimble reader should be used
   /// in this query or not.  Will be removed after the selective Nimble reader
   /// is fully rolled out.
diff --git a/presto-native-execution/presto_cpp/main/TaskManager.cpp b/presto-native-execution/presto_cpp/main/TaskManager.cpp
index 10ed292bac4f6..0b7fb7cb87b3c 100644
--- a/presto-native-execution/presto_cpp/main/TaskManager.cpp
+++ b/presto-native-execution/presto_cpp/main/TaskManager.cpp
@@ -536,13 +536,17 @@ std::unique_ptr TaskManager::createOrUpdateTaskImpl(
     bool summarize,
     std::shared_ptr queryCtx,
     long startProcessCpuTime) {
+  auto receiveTaskUpdateMs = getCurrentTimeMs();
   std::shared_ptr execTask;
   bool startTask = false;
   auto prestoTask = findOrCreateTask(taskId, startProcessCpuTime);
+  if (prestoTask->firstTimeReceiveTaskUpdateMs == 0) {
+    prestoTask->firstTimeReceiveTaskUpdateMs = receiveTaskUpdateMs;
+  }
   {
     std::lock_guard l(prestoTask->mutex);
     prestoTask->updateCoordinatorHeartbeatLocked();
-    if (not prestoTask->task && planFragment.planNode) {
+    if ((prestoTask->task == nullptr) && (planFragment.planNode != nullptr)) {
       // If the task is aborted, no need to do anything else.
       // This takes care of DELETE task message coming before CREATE task.
       if (prestoTask->info.taskStatus.state == protocol::TaskState::ABORTED) {
@@ -573,6 +577,7 @@ std::unique_ptr TaskManager::createOrUpdateTaskImpl(
       prestoTask->task = std::move(newExecTask);
       prestoTask->info.needsPlan = false;
       startTask = true;
+      prestoTask->createFinishTimeMs = getCurrentTimeMs();
     }
     execTask = prestoTask->task;
   }
@@ -610,7 +615,33 @@ std::unique_ptr TaskManager::createOrUpdateTaskImpl(
       VLOG(1) << "Failed to update output buffers for task: " << taskId;
     }
 
+    folly::F14FastMap sourcesMap;
     for (const auto& source : sources) {
+      auto it = sourcesMap.find(source.planNodeId);
+      if (it == sourcesMap.end()) {
+        // No existing source with same planNodeId, add as new
+        sourcesMap.emplace(source.planNodeId, source);
+        continue;
+      }
+
+      // Merge with existing source that has the same planNodeId
+      auto& merged = it->second;
+
+      // Merge splits
+      merged.splits.insert(
+          merged.splits.end(), source.splits.begin(), source.splits.end());
+
+      // Merge noMoreSplitsForLifespan
+      merged.noMoreSplitsForLifespan.insert(
+          merged.noMoreSplitsForLifespan.end(),
+          source.noMoreSplitsForLifespan.begin(),
+          source.noMoreSplitsForLifespan.end());
+
+      // Use OR logic for noMoreSplits flag
+      merged.noMoreSplits = merged.noMoreSplits || source.noMoreSplits;
+    }
+
+    for (const auto& [_, source] : sourcesMap) {
       // Add all splits from the source to the task.
       VLOG(1) << "Adding " << source.splits.size() << " splits to " << taskId
               << " for node " << source.planNodeId;
@@ -743,6 +774,8 @@ void TaskManager::startTaskLocked(std::shared_ptr& prestoTask) {
 
   // Record the time we spent between task creation and start, which is the
   // planned (queued) time.
+  // Note task could be created at getTaskStatus/getTaskInfo endpoint and later
+  // receive taskUpdate to create and start task.
   const auto queuedTimeInMs =
       velox::getCurrentTimeMs() - prestoTask->createTimeMs;
   prestoTask->info.stats.queuedTimeInNanos = queuedTimeInMs * 1'000'000;
diff --git a/presto-native-execution/presto_cpp/main/TaskResource.cpp b/presto-native-execution/presto_cpp/main/TaskResource.cpp
index 863ad0494925b..ec08cf19bcf02 100644
--- a/presto-native-execution/presto_cpp/main/TaskResource.cpp
+++ b/presto-native-execution/presto_cpp/main/TaskResource.cpp
@@ -231,9 +231,13 @@ proxygen::RequestHandler* TaskResource::createOrUpdateTaskImpl(
           std::shared_ptr handlerState) {
         folly::via(
             httpSrvCpuExecutor_,
-            [this, &body, taskId, summarize, createOrUpdateFunc, receiveThrift]() {
+            [this,
+             requestBody = util::extractMessageBody(body),
+             taskId,
+             summarize,
+             createOrUpdateFunc,
+             receiveThrift]() {
               const auto startProcessCpuTimeNs = util::getProcessCpuTimeNs();
-              std::string requestBody = util::extractMessageBody(body);
 
               std::unique_ptr taskInfo;
               try {
diff --git a/presto-native-execution/presto_cpp/main/common/Configs.cpp b/presto-native-execution/presto_cpp/main/common/Configs.cpp
index 5d8a86af6b596..76613ee440df4 100644
--- a/presto-native-execution/presto_cpp/main/common/Configs.cpp
+++ b/presto-native-execution/presto_cpp/main/common/Configs.cpp
@@ -261,6 +261,8 @@ SystemConfig::SystemConfig() {
           NUM_PROP(kHttpSrvIoEvbViolationThresholdMs, 1000),
           NUM_PROP(kMaxLocalExchangePartitionBufferSize, 65536),
           BOOL_PROP(kTextWriterEnabled, true),
+          BOOL_PROP(kCharNToVarcharImplicitCast, false),
+          BOOL_PROP(kEnumTypesEnabled, true),
       };
 }
 
@@ -926,6 +928,14 @@ bool SystemConfig::textWriterEnabled() const {
   return optionalProperty(kTextWriterEnabled).value();
 }
 
+bool SystemConfig::charNToVarcharImplicitCast() const {
+  return optionalProperty(kCharNToVarcharImplicitCast).value();
+}
+
+bool SystemConfig::enumTypesEnabled() const {
+  return optionalProperty(kEnumTypesEnabled).value();
+}
+
 NodeConfig::NodeConfig() {
   registeredProps_ =
       std::unordered_map>{
diff --git a/presto-native-execution/presto_cpp/main/common/Configs.h b/presto-native-execution/presto_cpp/main/common/Configs.h
index baaa9dbf06d4c..c891aee3d2306 100644
--- a/presto-native-execution/presto_cpp/main/common/Configs.h
+++ b/presto-native-execution/presto_cpp/main/common/Configs.h
@@ -763,6 +763,17 @@ class SystemConfig : public ConfigBase {
   // TODO: remove once text writer is fully rolled out
   static constexpr std::string_view kTextWriterEnabled{"text-writer-enabled"};
 
+  /// Enable the type char(n) with the same behavior as unbounded varchar.
+  /// char(n) type is not supported by parser when set to false.
+  static constexpr std::string_view kCharNToVarcharImplicitCast{
+    "char-n-to-varchar-implicit-cast"};
+
+  /// Enable BigintEnum and VarcharEnum types to be parsed and used in Velox. 
+  /// When set to false, BigintEnum or VarcharEnum types will throw an 
+  //  unsupported error during type parsing.
+  static constexpr std::string_view kEnumTypesEnabled{
+    "enum-types-enabled"};
+
   SystemConfig();
 
   virtual ~SystemConfig() = default;
@@ -1053,6 +1064,10 @@ class SystemConfig : public ConfigBase {
   uint64_t maxLocalExchangePartitionBufferSize() const;
 
   bool textWriterEnabled() const;
+
+  bool charNToVarcharImplicitCast() const;
+
+  bool enumTypesEnabled() const;
 };
 
 /// Provides access to node properties defined in node.properties file.
diff --git a/presto-native-execution/presto_cpp/main/common/tests/CMakeLists.txt b/presto-native-execution/presto_cpp/main/common/tests/CMakeLists.txt
index ea638403e6640..267118dd1cd2f 100644
--- a/presto-native-execution/presto_cpp/main/common/tests/CMakeLists.txt
+++ b/presto-native-execution/presto_cpp/main/common/tests/CMakeLists.txt
@@ -9,6 +9,13 @@
 # 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.
+add_library(presto_mutable_configs MutableConfigs.cpp)
+
+target_link_libraries(
+  presto_mutable_configs
+  presto_common
+  velox_file)
+
 add_executable(presto_common_test CommonTest.cpp ConfigTest.cpp)
 
 add_test(presto_common_test presto_common_test)
diff --git a/presto-native-execution/presto_cpp/main/tests/MutableConfigs.cpp b/presto-native-execution/presto_cpp/main/common/tests/MutableConfigs.cpp
similarity index 96%
rename from presto-native-execution/presto_cpp/main/tests/MutableConfigs.cpp
rename to presto-native-execution/presto_cpp/main/common/tests/MutableConfigs.cpp
index cd59c17b2ca97..1232cc0e4c5cf 100644
--- a/presto-native-execution/presto_cpp/main/tests/MutableConfigs.cpp
+++ b/presto-native-execution/presto_cpp/main/common/tests/MutableConfigs.cpp
@@ -12,7 +12,7 @@
  * limitations under the License.
  */
 #include "presto_cpp/main/common/Configs.h"
-#include "presto_cpp/main/tests/MultableConfigs.h"
+#include "presto_cpp/main/common/tests/MutableConfigs.h"
 #include "velox/common/file/File.h"
 #include "velox/common/file/FileSystems.h"
 #include "velox/exec/tests/utils/TempDirectoryPath.h"
diff --git a/presto-native-execution/presto_cpp/main/tests/MultableConfigs.h b/presto-native-execution/presto_cpp/main/common/tests/MutableConfigs.h
similarity index 100%
rename from presto-native-execution/presto_cpp/main/tests/MultableConfigs.h
rename to presto-native-execution/presto_cpp/main/common/tests/MutableConfigs.h
diff --git a/presto-native-execution/presto_cpp/main/connectors/SystemConnector.cpp b/presto-native-execution/presto_cpp/main/connectors/SystemConnector.cpp
index 2fc0b7338e351..d75dc5435c60b 100644
--- a/presto-native-execution/presto_cpp/main/connectors/SystemConnector.cpp
+++ b/presto-native-execution/presto_cpp/main/connectors/SystemConnector.cpp
@@ -74,6 +74,7 @@ SystemTableHandle::SystemTableHandle(
     std::string schemaName,
     std::string tableName)
     : ConnectorTableHandle(std::move(connectorId)),
+      name_(fmt::format("{}.{}", schemaName, tableName)),
       schemaName_(std::move(schemaName)),
       tableName_(std::move(tableName)) {
   VELOX_USER_CHECK_EQ(
diff --git a/presto-native-execution/presto_cpp/main/connectors/SystemConnector.h b/presto-native-execution/presto_cpp/main/connectors/SystemConnector.h
index 5e77af7a43c66..f491a4b3b98a9 100644
--- a/presto-native-execution/presto_cpp/main/connectors/SystemConnector.h
+++ b/presto-native-execution/presto_cpp/main/connectors/SystemConnector.h
@@ -42,6 +42,10 @@ class SystemTableHandle : public velox::connector::ConnectorTableHandle {
 
   std::string toString() const override;
 
+  const std::string& name() const override {
+    return name_;
+  }
+
   const std::string& schemaName() const {
     return schemaName_;
   }
@@ -53,6 +57,7 @@ class SystemTableHandle : public velox::connector::ConnectorTableHandle {
   const velox::RowTypePtr taskSchema() const;
 
  private:
+  const std::string name_;
   const std::string schemaName_;
   const std::string tableName_;
 };
diff --git a/presto-native-execution/presto_cpp/main/connectors/arrow_flight/ArrowFlightConfig.cpp b/presto-native-execution/presto_cpp/main/connectors/arrow_flight/ArrowFlightConfig.cpp
index c0900c6c41de4..18d637df02db8 100644
--- a/presto-native-execution/presto_cpp/main/connectors/arrow_flight/ArrowFlightConfig.cpp
+++ b/presto-native-execution/presto_cpp/main/connectors/arrow_flight/ArrowFlightConfig.cpp
@@ -43,4 +43,14 @@ std::optional ArrowFlightConfig::serverSslCertificate() const {
       config_->get(kServerSslCertificate));
 }
 
+std::optional ArrowFlightConfig::clientSslCertificate() const {
+  return static_cast>(
+      config_->get(kClientSslCertificate));
+}
+
+std::optional ArrowFlightConfig::clientSslKey() const {
+  return static_cast>(
+      config_->get(kClientSslKey));
+}
+
 } // namespace facebook::presto
diff --git a/presto-native-execution/presto_cpp/main/connectors/arrow_flight/ArrowFlightConfig.h b/presto-native-execution/presto_cpp/main/connectors/arrow_flight/ArrowFlightConfig.h
index 77ad8e9379cf3..4839733baaab0 100644
--- a/presto-native-execution/presto_cpp/main/connectors/arrow_flight/ArrowFlightConfig.h
+++ b/presto-native-execution/presto_cpp/main/connectors/arrow_flight/ArrowFlightConfig.h
@@ -38,6 +38,12 @@ class ArrowFlightConfig {
   static constexpr const char* kServerSslCertificate =
       "arrow-flight.server-ssl-certificate";
 
+  static constexpr const char* kClientSslCertificate =
+      "arrow-flight.client-ssl-certificate";
+
+  static constexpr const char* kClientSslKey =
+      "arrow-flight.client-ssl-key";
+
   std::string authenticatorName() const;
 
   std::optional defaultServerHostname() const;
@@ -50,6 +56,10 @@ class ArrowFlightConfig {
 
   std::optional serverSslCertificate() const;
 
+  std::optional clientSslCertificate() const;
+  
+  std::optional clientSslKey() const;
+
  private:
   const std::shared_ptr config_;
 };
diff --git a/presto-native-execution/presto_cpp/main/connectors/arrow_flight/ArrowFlightConnector.cpp b/presto-native-execution/presto_cpp/main/connectors/arrow_flight/ArrowFlightConnector.cpp
index 28e3e5b5f074c..5c02ac051e076 100644
--- a/presto-native-execution/presto_cpp/main/connectors/arrow_flight/ArrowFlightConnector.cpp
+++ b/presto-native-execution/presto_cpp/main/connectors/arrow_flight/ArrowFlightConnector.cpp
@@ -71,6 +71,26 @@ ArrowFlightConnector::initClientOpts(
     clientOpts->tls_root_certs = cert;
   }
 
+  auto clientCertPath = config->clientSslCertificate();
+  if (clientCertPath.has_value()) {
+    std::ifstream certFile(clientCertPath.value());
+    VELOX_CHECK(
+        certFile.is_open(), "Could not open client certificate at {}", clientCertPath.value());
+    clientOpts->cert_chain.assign(
+        (std::istreambuf_iterator(certFile)),
+        (std::istreambuf_iterator()));
+  }
+
+  auto clientKeyPath = config->clientSslKey();
+  if (clientKeyPath.has_value()) {
+    std::ifstream keyFile(clientKeyPath.value());
+    VELOX_CHECK(
+        keyFile.is_open(), "Could not open client key at {}", clientKeyPath.value());
+    clientOpts->private_key.assign(
+        (std::istreambuf_iterator(keyFile)),
+        (std::istreambuf_iterator()));
+  }
+
   return clientOpts;
 }
 
diff --git a/presto-native-execution/presto_cpp/main/connectors/arrow_flight/ArrowFlightConnector.h b/presto-native-execution/presto_cpp/main/connectors/arrow_flight/ArrowFlightConnector.h
index 6123b270e2349..d1c431575c677 100644
--- a/presto-native-execution/presto_cpp/main/connectors/arrow_flight/ArrowFlightConnector.h
+++ b/presto-native-execution/presto_cpp/main/connectors/arrow_flight/ArrowFlightConnector.h
@@ -31,7 +31,14 @@ namespace facebook::presto {
 class ArrowFlightTableHandle : public velox::connector::ConnectorTableHandle {
  public:
   explicit ArrowFlightTableHandle(const std::string& connectorId)
-      : ConnectorTableHandle(connectorId) {}
+      : ConnectorTableHandle(connectorId), name_("arrow_flight") {}
+
+  const std::string& name() const override {
+    return name_;
+  }
+
+ private:
+  const std::string name_;
 };
 
 struct ArrowFlightSplit : public velox::connector::ConnectorSplit {
diff --git a/presto-native-execution/presto_cpp/main/connectors/arrow_flight/tests/ArrowFlightConfigTest.cpp b/presto-native-execution/presto_cpp/main/connectors/arrow_flight/tests/ArrowFlightConfigTest.cpp
index eb946f1fcae76..8ec912b417650 100644
--- a/presto-native-execution/presto_cpp/main/connectors/arrow_flight/tests/ArrowFlightConfigTest.cpp
+++ b/presto-native-execution/presto_cpp/main/connectors/arrow_flight/tests/ArrowFlightConfigTest.cpp
@@ -27,6 +27,8 @@ TEST(ArrowFlightConfigTest, defaultConfig) {
   ASSERT_EQ(config.defaultServerSslEnabled(), false);
   ASSERT_EQ(config.serverVerify(), true);
   ASSERT_EQ(config.serverSslCertificate(), std::nullopt);
+  ASSERT_EQ(config.clientSslCertificate(), std::nullopt);
+  ASSERT_EQ(config.clientSslKey(), std::nullopt);
 }
 
 TEST(ArrowFlightConfigTest, overrideConfig) {
@@ -36,7 +38,9 @@ TEST(ArrowFlightConfigTest, overrideConfig) {
       {ArrowFlightConfig::kDefaultServerPort, "9000"},
       {ArrowFlightConfig::kDefaultServerSslEnabled, "true"},
       {ArrowFlightConfig::kServerVerify, "false"},
-      {ArrowFlightConfig::kServerSslCertificate, "my-cert.crt"}};
+      {ArrowFlightConfig::kServerSslCertificate, "my-cert.crt"},
+      {ArrowFlightConfig::kClientSslCertificate, "/path/to/client.crt"},
+      {ArrowFlightConfig::kClientSslKey, "/path/to/client.key"}};
   auto config = ArrowFlightConfig(
       std::make_shared(std::move(configMap)));
   ASSERT_EQ(config.authenticatorName(), "my-authenticator");
@@ -45,4 +49,6 @@ TEST(ArrowFlightConfigTest, overrideConfig) {
   ASSERT_EQ(config.defaultServerSslEnabled(), true);
   ASSERT_EQ(config.serverVerify(), false);
   ASSERT_EQ(config.serverSslCertificate(), "my-cert.crt");
+  ASSERT_EQ(config.clientSslCertificate(), "/path/to/client.crt");
+  ASSERT_EQ(config.clientSslKey(), "/path/to/client.key");
 }
diff --git a/presto-native-execution/presto_cpp/main/connectors/arrow_flight/tests/ArrowFlightConnectorMTlsTest.cpp b/presto-native-execution/presto_cpp/main/connectors/arrow_flight/tests/ArrowFlightConnectorMTlsTest.cpp
new file mode 100644
index 0000000000000..8225172c8b526
--- /dev/null
+++ b/presto-native-execution/presto_cpp/main/connectors/arrow_flight/tests/ArrowFlightConnectorMTlsTest.cpp
@@ -0,0 +1,158 @@
+/*
+ * Licensed 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.
+ */
+#include 
+#include 
+#include 
+#include 
+#include 
+#include "presto_cpp/main/connectors/arrow_flight/ArrowFlightConnector.h"
+#include "presto_cpp/main/connectors/arrow_flight/tests/utils/ArrowFlightConnectorTestBase.h"
+#include "presto_cpp/main/connectors/arrow_flight/tests/utils/ArrowFlightPlanBuilder.h"
+#include "presto_cpp/main/connectors/arrow_flight/tests/utils/Utils.h"
+#include "velox/common/base/tests/GTestUtils.h"
+#include "velox/common/config/Config.h"
+#include "velox/exec/tests/utils/AssertQueryBuilder.h"
+
+using namespace arrow;
+using namespace facebook::velox;
+using namespace facebook::velox::exec::test;
+
+namespace facebook::presto::test {
+
+class ArrowFlightConnectorMtlsTestBase : public ArrowFlightConnectorTestBase {
+ protected:
+  explicit ArrowFlightConnectorMtlsTestBase(
+      std::shared_ptr config)
+      : ArrowFlightConnectorTestBase(std::move(config)) {}
+
+  void setFlightServerOptions(
+      flight::FlightServerOptions* serverOptions) override {
+    flight::CertKeyPair tlsCertificate{
+        .pem_cert = readFile("./data/certs/server.crt"),
+        .pem_key = readFile("./data/certs/server.key")};
+    serverOptions->tls_certificates.push_back(tlsCertificate);
+    serverOptions->verify_client = true;
+    serverOptions->root_certificates = readFile("./data/certs/ca.crt");
+  }
+
+  void executeSuccessfulQuery() {
+    std::vector idData = {
+        1, 12, 2, std::numeric_limits::max()};
+
+    updateTable(
+        "sample-data",
+        makeArrowTable({"id"}, {makeNumericArray(idData)}));
+
+    auto idVec = makeFlatVector(idData);
+
+    auto plan = ArrowFlightPlanBuilder()
+                    .flightTableScan(ROW({"id"}, {BIGINT()}))
+                    .planNode();
+
+    AssertQueryBuilder(plan)
+        .splits(makeSplits({"sample-data"}))
+        .assertResults(makeRowVector({idVec}));
+  }
+
+  std::function createQueryFunction() {
+    std::vector idData = {
+        1, 12, 2, std::numeric_limits::max()};
+
+    updateTable(
+        "sample-data",
+        makeArrowTable({"id"}, {makeNumericArray(idData)}));
+
+    auto idVec = makeFlatVector(idData);
+
+    auto plan = ArrowFlightPlanBuilder()
+                    .flightTableScan(ROW({"id"}, {BIGINT()}))
+                    .planNode();
+
+    return [this, plan, idVec]() {
+      AssertQueryBuilder(plan)
+          .splits(makeSplits({"sample-data"}))
+          .assertResults(makeRowVector({idVec}));
+    };
+  }
+};
+
+class ArrowFlightConnectorMtlsTest : public ArrowFlightConnectorMtlsTestBase {
+ protected:
+  explicit ArrowFlightConnectorMtlsTest()
+      : ArrowFlightConnectorMtlsTestBase(
+            std::make_shared(
+                std::unordered_map{
+                    {ArrowFlightConfig::kDefaultServerSslEnabled, "true"},
+                    {ArrowFlightConfig::kServerVerify, "true"},
+                    {ArrowFlightConfig::kServerSslCertificate,
+                         "./data/certs/ca.crt"},
+                    {ArrowFlightConfig::kClientSslCertificate, "./data/certs/client.crt"},
+                    {ArrowFlightConfig::kClientSslKey, "./data/certs/client.key"}})) {}
+};
+
+TEST_F(ArrowFlightConnectorMtlsTest, successfulMtlsConnection) {
+  executeSuccessfulQuery();
+}
+
+class ArrowFlightMtlsNoClientCertTest : public ArrowFlightConnectorMtlsTestBase {
+ protected:
+  ArrowFlightMtlsNoClientCertTest()
+      : ArrowFlightConnectorMtlsTestBase(
+            std::make_shared(
+                std::unordered_map{
+                    {ArrowFlightConfig::kDefaultServerSslEnabled, "true"},
+                    {ArrowFlightConfig::kServerVerify, "true"},
+                    {ArrowFlightConfig::kServerSslCertificate, "./data/certs/ca.crt"}})) {}
+};
+
+TEST_F(ArrowFlightMtlsNoClientCertTest, mtlsFailsWithoutClientCert) {
+  auto queryFunction = createQueryFunction();
+  VELOX_ASSERT_THROW(queryFunction(), "failed to connect");
+}
+
+class ArrowFlightConnectorImplicitSslTest : public ArrowFlightConnectorMtlsTestBase {
+ protected:
+  ArrowFlightConnectorImplicitSslTest()
+      : ArrowFlightConnectorMtlsTestBase(
+            std::make_shared(
+                std::unordered_map{
+                    {ArrowFlightConfig::kServerVerify, "true"},
+                    {ArrowFlightConfig::kServerSslCertificate,
+                         "./data/certs/ca.crt"},
+                    {ArrowFlightConfig::kClientSslCertificate, "./data/certs/client.crt"},
+                    {ArrowFlightConfig::kClientSslKey, "./data/certs/client.key"}})) {}
+};
+
+TEST_F(ArrowFlightConnectorImplicitSslTest, successfulImplicitSslConnection) {
+  executeSuccessfulQuery();
+}
+
+class ArrowFlightImplicitSslNoClientCertTest : public ArrowFlightConnectorMtlsTestBase {
+ protected:
+  ArrowFlightImplicitSslNoClientCertTest()
+      : ArrowFlightConnectorMtlsTestBase(
+            std::make_shared(
+                std::unordered_map{
+                    {ArrowFlightConfig::kDefaultServerSslEnabled, "true"},
+                    {ArrowFlightConfig::kServerVerify, "true"},
+                    {ArrowFlightConfig::kServerSslCertificate, "./data/certs/ca.crt"}
+                    })) {}
+};
+
+TEST_F(ArrowFlightImplicitSslNoClientCertTest, mtlsFailsWithoutClientCertOnImplicitSsl) {
+  auto queryFunction = createQueryFunction();
+  VELOX_ASSERT_THROW(queryFunction(), "failed to connect");
+}
+
+} // namespace facebook::presto::test
diff --git a/presto-native-execution/presto_cpp/main/connectors/arrow_flight/tests/ArrowFlightConnectorTlsTest.cpp b/presto-native-execution/presto_cpp/main/connectors/arrow_flight/tests/ArrowFlightConnectorTlsTest.cpp
index 4453183a39412..0be7df437e5e4 100644
--- a/presto-native-execution/presto_cpp/main/connectors/arrow_flight/tests/ArrowFlightConnectorTlsTest.cpp
+++ b/presto-native-execution/presto_cpp/main/connectors/arrow_flight/tests/ArrowFlightConnectorTlsTest.cpp
@@ -39,8 +39,8 @@ class ArrowFlightConnectorTlsTestBase : public ArrowFlightConnectorTestBase {
   void setFlightServerOptions(
       flight::FlightServerOptions* serverOptions) override {
     flight::CertKeyPair tlsCertificate{
-        .pem_cert = readFile("./data/tls_certs/server.crt"),
-        .pem_key = readFile("./data/tls_certs/server.key")};
+        .pem_cert = readFile("./data/certs/server.crt"),
+        .pem_key = readFile("./data/certs/server.key")};
     serverOptions->tls_certificates.push_back(tlsCertificate);
   }
 
@@ -83,7 +83,7 @@ class ArrowFlightConnectorTlsTest : public ArrowFlightConnectorTlsTestBase {
                     {ArrowFlightConfig::kDefaultServerSslEnabled, "true"},
                     {ArrowFlightConfig::kServerVerify, "true"},
                     {ArrowFlightConfig::kServerSslCertificate,
-                     "./data/tls_certs/ca.crt"}})) {}
+                     "./data/certs/ca.crt"}})) {}
 };
 
 TEST_F(ArrowFlightConnectorTlsTest, tlsEnabled) {
@@ -116,7 +116,7 @@ class ArrowFlightTlsNoCertTest : public ArrowFlightConnectorTlsTestBase {
 };
 
 TEST_F(ArrowFlightTlsNoCertTest, tlsNoCert) {
-  executeTest(false, "handshake failed");
+  executeTest(false, "failed to connect");
 }
 
 } // namespace facebook::presto::test
diff --git a/presto-native-execution/presto_cpp/main/connectors/arrow_flight/tests/CMakeLists.txt b/presto-native-execution/presto_cpp/main/connectors/arrow_flight/tests/CMakeLists.txt
index 9af596a913973..f346d1212a1ad 100644
--- a/presto-native-execution/presto_cpp/main/connectors/arrow_flight/tests/CMakeLists.txt
+++ b/presto-native-execution/presto_cpp/main/connectors/arrow_flight/tests/CMakeLists.txt
@@ -23,16 +23,19 @@ target_link_libraries(
 add_executable(
   presto_flight_connector_test
   ArrowFlightConnectorTest.cpp ArrowFlightConnectorAuthTest.cpp
+  ArrowFlightConnectorMTlsTest.cpp
   ArrowFlightConnectorTlsTest.cpp ArrowFlightConnectorDataTypeTest.cpp
   ArrowFlightConfigTest.cpp)
 
-set(DATA_DIR "${CMAKE_CURRENT_SOURCE_DIR}/data/tls_certs")
-
 add_custom_target(
   copy_flight_test_data ALL
-  COMMAND ${CMAKE_COMMAND} -E copy_directory ${DATA_DIR}
-          $/data/tls_certs)
+  COMMAND ${CMAKE_COMMAND} -E copy_directory_if_different
+          "${CMAKE_CURRENT_SOURCE_DIR}/data"
+          "${CMAKE_CURRENT_BINARY_DIR}/data"
+  COMMENT "Copying test data files..."
+)
 
+add_dependencies(presto_flight_connector_test copy_flight_test_data)
 add_test(presto_flight_connector_test presto_flight_connector_test)
 
 target_link_libraries(
diff --git a/presto-native-execution/presto_cpp/main/connectors/arrow_flight/tests/data/README.md b/presto-native-execution/presto_cpp/main/connectors/arrow_flight/tests/data/README.md
index 3a5f2e5786c67..17e778a84750c 100644
--- a/presto-native-execution/presto_cpp/main/connectors/arrow_flight/tests/data/README.md
+++ b/presto-native-execution/presto_cpp/main/connectors/arrow_flight/tests/data/README.md
@@ -1,7 +1,7 @@
-### Placeholder TLS Certificates for Arrow Flight Connector Unit Testing
-The `tls_certs` directory contains placeholder TLS certificates generated for unit testing the Arrow Flight Connector with TLS enabled. These certificates are not intended for production use and should only be used in the context of unit tests.
+### Placeholder TLS & mTLS Certificates for Arrow Flight Connector Unit Testing
+The `certs/` directory contains placeholder certificates used for unit and integration testing of the Arrow Flight Connector with **TLS** and **mutual TLS (mTLS)** enabled. These certificates are **not intended for production use** and should only be used for local development or automated testing scenarios.
 
-### Generating TLS Certificates
-To create the TLS certificates and keys inside the `tls_certs` folder, run the following command:
+### Generating Certificates for TLS and mTLS Testing
+To generate the necessary certificates and keys for both TLS and mTLS testing, run the following script:
 
-`./generate_tls_certs.sh`
+`./generate_certs.sh`
diff --git a/presto-native-execution/presto_cpp/main/connectors/arrow_flight/tests/data/certs/ca.crt b/presto-native-execution/presto_cpp/main/connectors/arrow_flight/tests/data/certs/ca.crt
new file mode 100644
index 0000000000000..bc6036c52aa12
--- /dev/null
+++ b/presto-native-execution/presto_cpp/main/connectors/arrow_flight/tests/data/certs/ca.crt
@@ -0,0 +1,22 @@
+-----BEGIN CERTIFICATE-----
+MIIDmzCCAoOgAwIBAgIUVB6MDVAXGLccJY5XrssVNLKbH/gwDQYJKoZIhvcNAQEL
+BQAwXDELMAkGA1UEBhMCVVMxDjAMBgNVBAgMBVN0YXRlMQ0wCwYDVQQHDARDaXR5
+MQ4wDAYDVQQKDAVNeU9yZzEPMA0GA1UECwwGTXlVbml0MQ0wCwYDVQQDDARNeUNB
+MCAXDTI1MDczMDE2NTU1M1oYDzMwMjQxMTMwMTY1NTUzWjBcMQswCQYDVQQGEwJV
+UzEOMAwGA1UECAwFU3RhdGUxDTALBgNVBAcMBENpdHkxDjAMBgNVBAoMBU15T3Jn
+MQ8wDQYDVQQLDAZNeVVuaXQxDTALBgNVBAMMBE15Q0EwggEiMA0GCSqGSIb3DQEB
+AQUAA4IBDwAwggEKAoIBAQCQ5D8+R5/HesO9BrgWF92wGOesBcOPQL7Nemzh1qYT
+fhxNV7HGN1W5fByWsVelZ9326CTa7Yr4TlRXC8GHt9YbRvU68LBU1kAGUqGmPdHY
+LgPZJkQqFhYdPdD++Gw2+ecfYp+Ls9pZm/pNOeCWBV0F3RaoPZ19m/C/lDdZz8OX
+V6t2to+Yh0GXJyyJfO+w7qG6B/j8UiRYtnnMq3ywTcWqsUmYp4+uwmkEEN9/eSo9
+JanjpEiv6o9Yb9J5StXRPmbAoXOl45o87A6qo0vzgYdP6uKkPQxo6wSdltb5qctM
+CKHm4bYFT6IHoVzUGVHiq2iRna87OiDPqnFaX3ktWJB9AgMBAAGjUzBRMB0GA1Ud
+DgQWBBRGccIvJxNjWYY8hGZ3HIlut0p17jAfBgNVHSMEGDAWgBRGccIvJxNjWYY8
+hGZ3HIlut0p17jAPBgNVHRMBAf8EBTADAQH/MA0GCSqGSIb3DQEBCwUAA4IBAQAI
+sUZOfL8ZJEu0Oq7AHgOmxMYhdRQmFr45C54fiNb6QLs5bPJZ9T1/nbPQZvF1hkFh
+wB3pDn8rNntiYXkCkrQ6PAiQfn0WHl4jLYCoGYxFbSP1QViZNid7dPmpaxccjMhL
+Zk7htfCS1HtHWWBZPMDDA8hsUvBf4qusVonO71XGL22Z2ZKtgvDJYAyoxm7xwIo2
+mqSH9TfOnHYE0hUpo3u4PdmVAfCXzSDRccLALnVlzt50ColmAQgzj3MnwWfXJmdv
+kjBhIZ9Obt7Hf6FcBhX9/qQN1t3u6mLIjf2akokRFblcW5Xgv4X1c7pEtdtM9V+o
+MVqiDfm49mpMcOdHDeb1
+-----END CERTIFICATE-----
diff --git a/presto-native-execution/presto_cpp/main/connectors/arrow_flight/tests/data/certs/client.crt b/presto-native-execution/presto_cpp/main/connectors/arrow_flight/tests/data/certs/client.crt
new file mode 100644
index 0000000000000..7313fbf6babc6
--- /dev/null
+++ b/presto-native-execution/presto_cpp/main/connectors/arrow_flight/tests/data/certs/client.crt
@@ -0,0 +1,22 @@
+-----BEGIN CERTIFICATE-----
+MIIDmTCCAoGgAwIBAgIUHZmKZs2+ejbJGEezgoSUTPKF42YwDQYJKoZIhvcNAQEL
+BQAwXDELMAkGA1UEBhMCVVMxDjAMBgNVBAgMBVN0YXRlMQ0wCwYDVQQHDARDaXR5
+MQ4wDAYDVQQKDAVNeU9yZzEPMA0GA1UECwwGTXlVbml0MQ0wCwYDVQQDDARNeUNB
+MCAXDTI1MDczMDE2NTU1M1oYDzMwMjQxMTMwMTY1NTUzWjBrMQswCQYDVQQGEwJV
+UzEOMAwGA1UECAwFU3RhdGUxDTALBgNVBAcMBENpdHkxDjAMBgNVBAoMBU15T3Jn
+MQ8wDQYDVQQLDAZNeVVuaXQxHDAaBgNVBAMME2NsaWVudC5teWRvbWFpbi5jb20w
+ggEiMA0GCSqGSIb3DQEBAQUAA4IBDwAwggEKAoIBAQDESuInkBMIwJ02rNuO1DPN
+3BfuNzF2SJq5/UgH9soWxVIYqLOwinBmWRsrRVOps9mJjN/GYbj27GT1LFXZHIwh
+MYf5JLE/Y8CL/JwnMyxe3wkiBb1s2d+urOrIPMLQTVLGtClyPVT1PnmOd/h7vJ+k
+I5ciK9G+8krAq5NoWM/OvhdTK2pOns1jl6Lq2c4IGDgQHt65uKRqPTPtc1xe/Vxd
+4sTT8LJQEHu6zdl5nGFdGq4eIJpr+9LuM8ZCYkTAYpzHJ2XKP7jEj6ZRmsIgwwTf
+VMON+yAF06Tmk9VCRkLUGjNMA4Xh24MejieE0JSq7mIAhSYWE+nalG8HsxtsQWDt
+AgMBAAGjQjBAMB0GA1UdDgQWBBRPc2AWGMwuJEbnnMXvzgNAizik1jAfBgNVHSME
+GDAWgBRGccIvJxNjWYY8hGZ3HIlut0p17jANBgkqhkiG9w0BAQsFAAOCAQEACJde
+pXZvE04uw6tv+iGplmYNfasMQ9JXbvi0JMlnt9Y7ajf0F3g6yw9xWQfp8mCWVuof
+aS/X8qw3loundeprxVq/2V6pFXStLFXJCXX+YL0Wl8AMv8VOxdZ8+hYlkfMoiKnx
+ZKaWgrVtI/idFRUJLg8aHLRk5qVOwACBg9DAxMBC4V4MCQBfvDDdY6Y5qAM8o7PN
+YGsPv5JIQI/3jsG2ZNQ/A8Ar+BNKWqnwRg2jjXysjPJPaU8TExvFXUypQaUn+a8X
+/Y+CXhGabfbodEbEvBny6tiFQMb3YWhd0kHkjrYylY8GOpQ/ziEC8s7JYcaLSiu/
+ko1ErmRY/vhQIvrylw==
+-----END CERTIFICATE-----
diff --git a/presto-native-execution/presto_cpp/main/connectors/arrow_flight/tests/data/certs/client.key b/presto-native-execution/presto_cpp/main/connectors/arrow_flight/tests/data/certs/client.key
new file mode 100644
index 0000000000000..2c285f776ba2d
--- /dev/null
+++ b/presto-native-execution/presto_cpp/main/connectors/arrow_flight/tests/data/certs/client.key
@@ -0,0 +1,28 @@
+-----BEGIN PRIVATE KEY-----
+MIIEvQIBADANBgkqhkiG9w0BAQEFAASCBKcwggSjAgEAAoIBAQDESuInkBMIwJ02
+rNuO1DPN3BfuNzF2SJq5/UgH9soWxVIYqLOwinBmWRsrRVOps9mJjN/GYbj27GT1
+LFXZHIwhMYf5JLE/Y8CL/JwnMyxe3wkiBb1s2d+urOrIPMLQTVLGtClyPVT1PnmO
+d/h7vJ+kI5ciK9G+8krAq5NoWM/OvhdTK2pOns1jl6Lq2c4IGDgQHt65uKRqPTPt
+c1xe/Vxd4sTT8LJQEHu6zdl5nGFdGq4eIJpr+9LuM8ZCYkTAYpzHJ2XKP7jEj6ZR
+msIgwwTfVMON+yAF06Tmk9VCRkLUGjNMA4Xh24MejieE0JSq7mIAhSYWE+nalG8H
+sxtsQWDtAgMBAAECggEAKjnqszSS56nZ2Bpw4+Wh3EnZywjMBuRBBrwmC/KK0EGz
+8rKN7y8k1VubVOBlykayiBzKQciha9r4L+bQ8/LocTaQx+edCqQolmSp6ePgCmuj
+8RH3iSxIanDv09IAXaOYqD63AMiRV22QZDXKOkIePIbcewEerps8Ofze6c5bK9/H
+jWWScpphiJcTVSNKqTzVMJ27DcpTn/bQ+nPYQNdyMj51d1Vaebesmv1ieQfvdMTA
+FFTCYUfNrWWjB9CEhZAg3sgr3dOVVTl8LC8VF6ug5EvNeuW2l8FX+fgVaZWKzGpf
+veItiHIeMAEw6ZQPTumSvIO2fhg5Iue9fu4hR08twQKBgQDuqsRTJNG8DhR2CZrH
+KWeW8ac7pOiHSHxqpERfwkxtpUEsLXib8foPRJYdEzvDDgL3SxaRmTnwtyZvhZgl
+RE7ebRvAHkHy+nvim6wBgFYe2YwTTm9d3gCRxOHknEXfXSxzbwJznDlZrHlJ9XDQ
+OFBgLbhheF694y1l605E/8s/QQKBgQDSjEztXmrzlQvXKhOp+eL6o+jRrfrv5hmC
+ORy063ecVD/bA3OsgOKd5Q0v7dEyCC5IesQA0uY4GZ7zjsxcA4tUp20fZoTEeMpZ
+DznTz5Lqkpa8DCeMLj6fPWKnb4zZeduvBE4BRzS4Mdpzp/OrLAGgE2hwWNIB56uq
+pyloAkEirQKBgQDPZjg7JFjaOcYQGSKWheWOJyszSogC37u2lE8Sg+8UrTGoaU9Q
+/QNXdzuXwpoBU9DCA092cRgHlbDh4s8nO2fqJBikZ+bZdlBnyO29VEACiPwP3u4q
+PPxzsAq5NhAGHZq+KS6RNqYjxhyUZ6SEXRuDqNd8ZDS4gI137vZSQZLmwQKBgDcs
+mQQjF/fY+Q9bcWe7miWASoSYCQhQziJ4APPQOLn4wfsMvoVYCQrDeV8z/PwVdLt9
+oFtu6PGOlT7SDu+V5i866Levz98EoFISUV8WKDPcUi/ZJ4vumm50UaP68XgUHOOS
+RzbCiCg0uEBSpOIYWBywuU+nlvD02uGPiKQ+4v7JAoGAdyREhzfyz+WO5wdHLRIB
+k9cz448Lp/Uh2pcgAkFvtBwCAGZn86YJZ7JFD/HkQbw7amVEn+nxoCeWHkbSnzt0
+8Gu0hdqo6SaxOdf1wKBel79r/GH4ZwYRzWJoO/Iy//JJJO7g9GylAo0rBdQR9n60
+ZPMIZJLfKmGQWahl2EqyMG8=
+-----END PRIVATE KEY-----
diff --git a/presto-native-execution/presto_cpp/main/connectors/arrow_flight/tests/data/certs/server.crt b/presto-native-execution/presto_cpp/main/connectors/arrow_flight/tests/data/certs/server.crt
new file mode 100644
index 0000000000000..bdae53b6bd114
--- /dev/null
+++ b/presto-native-execution/presto_cpp/main/connectors/arrow_flight/tests/data/certs/server.crt
@@ -0,0 +1,23 @@
+-----BEGIN CERTIFICATE-----
+MIIDxDCCAqygAwIBAgIUHZmKZs2+ejbJGEezgoSUTPKF42UwDQYJKoZIhvcNAQEL
+BQAwXDELMAkGA1UEBhMCVVMxDjAMBgNVBAgMBVN0YXRlMQ0wCwYDVQQHDARDaXR5
+MQ4wDAYDVQQKDAVNeU9yZzEPMA0GA1UECwwGTXlVbml0MQ0wCwYDVQQDDARNeUNB
+MCAXDTI1MDczMDE2NTU1M1oYDzMwMjQxMTMwMTY1NTUzWjBrMQswCQYDVQQGEwJV
+UzEOMAwGA1UECAwFU3RhdGUxDTALBgNVBAcMBENpdHkxDjAMBgNVBAoMBU15T3Jn
+MQ8wDQYDVQQLDAZNeVVuaXQxHDAaBgNVBAMME3NlcnZlci5teWRvbWFpbi5jb20w
+ggEiMA0GCSqGSIb3DQEBAQUAA4IBDwAwggEKAoIBAQDS9+yttKuydwcc4te8CtJS
+yldr163zuSemazYUDY07ZMc/VFfSISYBX3yrW8uveR4zt8H8XhdkPHV6t3K3GxcY
+28CpUXr9TkCH2xMd4xn2DuewaTuf+yCcM6TLFh00nyqcYPhZGn3Bd5jGXKxI8PsH
+butswFKM/t9VTRtTpBgMsw8SEy0vyvsJyPTHm8aWY7tCDSI+vI8bHcG9sO8cuin8
+0JcZ4rRTgZpmDDlcY0OniRwowB5ph3eU0uaHIX+EWuht3+1trjyrSFuQ6y82f/0N
+MchcA3vTG0qF4ulcZn0yng1wuC1YorInYiDUahxDXHvmsoYz+0yOkR5fVReHLfdp
+AgMBAAGjbTBrMCkGA1UdEQQiMCCCE3NlcnZlci5teWRvbWFpbi5jb22CCWxvY2Fs
+aG9zdDAdBgNVHQ4EFgQUuxBCAwXgk1cEC3myPm4Q/b8+oWMwHwYDVR0jBBgwFoAU
+RnHCLycTY1mGPIRmdxyJbrdKde4wDQYJKoZIhvcNAQELBQADggEBABlwupm4YIBd
+dPRoX6S/Ta8oAvkz+gv/s9vIekG7fTzfGDmh9hQEtVU4OVT0ifVAX7x5bf0v7KKp
+jXDctZqyttXAtu6e1nQBE+MC9MXi68YU2hSApJF+z7WTPXEbcuIOQPpvXfiD7s2j
+455tMF90iKWjcFGdgB1usiQeNeNDjBwcvkGJhLhKO6UqsHLh2BBUTPhDnpyBkt1y
+zsZk3YDKr2ipaIj3OvBwueJNL8I3oU04eZMAFeSYbnJ49GdSDLTxglACoOkZrbju
+4qwV1bbIB97AtSdhcg3OxBZTAgVep8kOl3N/spNfgq4N/jTdsl+hoEKlChH/7kI3
+oYZfv+nc6s0=
+-----END CERTIFICATE-----
diff --git a/presto-native-execution/presto_cpp/main/connectors/arrow_flight/tests/data/certs/server.key b/presto-native-execution/presto_cpp/main/connectors/arrow_flight/tests/data/certs/server.key
new file mode 100644
index 0000000000000..1906728090bb7
--- /dev/null
+++ b/presto-native-execution/presto_cpp/main/connectors/arrow_flight/tests/data/certs/server.key
@@ -0,0 +1,28 @@
+-----BEGIN PRIVATE KEY-----
+MIIEvQIBADANBgkqhkiG9w0BAQEFAASCBKcwggSjAgEAAoIBAQDS9+yttKuydwcc
+4te8CtJSyldr163zuSemazYUDY07ZMc/VFfSISYBX3yrW8uveR4zt8H8XhdkPHV6
+t3K3GxcY28CpUXr9TkCH2xMd4xn2DuewaTuf+yCcM6TLFh00nyqcYPhZGn3Bd5jG
+XKxI8PsHbutswFKM/t9VTRtTpBgMsw8SEy0vyvsJyPTHm8aWY7tCDSI+vI8bHcG9
+sO8cuin80JcZ4rRTgZpmDDlcY0OniRwowB5ph3eU0uaHIX+EWuht3+1trjyrSFuQ
+6y82f/0NMchcA3vTG0qF4ulcZn0yng1wuC1YorInYiDUahxDXHvmsoYz+0yOkR5f
+VReHLfdpAgMBAAECggEAGOjA/zmH1EiNhHGcO02jy7asX8VVeqNv9QxPlEqNVGfv
+xqB0xhC35g2aMLlj8VIBqOWXd+68IE+rJ1QlrUz7iynXM6a1ONdWczQAq9S2qgDU
+hlXGfnsuPIM0f+4agK4SX+hrKkogcwll9nXWub4KRbRpA6wpkxA82luCUHvdgxIi
+PW37HHYWUKoAu828PpG1wGUf6wDFnSEhEuYfgHGkcWyhpRw5QZP563QTMCKWPjcW
+Nm+XJ10gpAB4Q+zcqxji7r+5uIbF9Zobkd9VuaWHavQfrgHpjfLnZjHr3r3BABi0
+U1Zz5x5R19r54aY93J0Fq4DOlJ1Gf80eBLJEs1b8CQKBgQD53IfUp4ScfIs1I3Vw
+JOclRY3UY/uZDQ+edXLqT91UxRpmvq8pmL82Y+idWGrx1WtCk2uOFGAl37Y+8goW
+bMHEZ8Wv8NtV29sqduuCE210miZN3q7EmTEj0AOOKd/Skwyte3rxvFw0didQvAhE
++uZEIZ+XaUNF102hT2BaTIm8fwKBgQDYJsgiqGRzGtP3sOeLKxYmGDE3fGB4JHYU
+U9kZ0pTQWiYsl+F7lXdwqkUApgU1rFuA3oR7dV4a8zS7BbyLK2WYfWE4yAoRm52W
+VnIGsdG4Z4wFGuNR7d+m+MouP4HYSJFtUoJFJxYXU4Kc3H88Ob5UmerVNFN7VSaU
+W/jtek34FwKBgH+c4sL5zAEgmvjI43IjZuriW03ewuGoihGkaszBfYmOIa3YNh5I
+pWBiJqw2PGjHV8DpCkXGolS1rZ74f650XYKyfYUevudbItTNZ/tHcN/c2zNqSFig
+5TglRauWN3qVICR6rJBKY81niyzw3Ehe3Lxvb9MlL/a7wCpjIBL+hFqBAoGBAIax
+jA+EzauosSPtWiwv+kpc0vaXi+nyFp7OLUBZKCC5vIYXUwxW9KoBgKRJ0H9E23Rv
+tTDVz4GNwnM0vOwga9vdbaMbjKKyTT4sujuPvXdjFy7rNXKNf8wlxp+RNZGYjv8H
+5mO/WpXIlWC4SpU2CnPfwiV/yPHW+waCVZlumH2bAoGAJscnMS4YQXiSoLQweiyQ
++mev9+i9h57/M0bk6WQHWzhAjxmZDtFUPO8P4vi3akAhCnt0yswh8e3G6h3F9TEz
+PQ4AIFMajb0PWUEcGsZhBh7TCJwxPXo34SbXGkIB5RYogT/mDVMjEfduAMen/0sp
+Te0BRozzkKHtIcGAqozt7ZU=
+-----END PRIVATE KEY-----
diff --git a/presto-native-execution/presto_cpp/main/connectors/arrow_flight/tests/data/generate_certs.sh b/presto-native-execution/presto_cpp/main/connectors/arrow_flight/tests/data/generate_certs.sh
new file mode 100755
index 0000000000000..11875d6fe462a
--- /dev/null
+++ b/presto-native-execution/presto_cpp/main/connectors/arrow_flight/tests/data/generate_certs.sh
@@ -0,0 +1,55 @@
+#!/bin/bash
+
+# Set directory for certificates and keys.
+CERT_DIR="./certs"
+mkdir -p $CERT_DIR
+
+# Dummy values for the certificates.
+COUNTRY="US"
+STATE="State"
+LOCALITY="City"
+ORGANIZATION="MyOrg"
+ORG_UNIT="MyUnit"
+COMMON_NAME="MyCA"
+SERVER_CN="server.mydomain.com"
+CLIENT_CN="client.mydomain.com"
+
+# Step 1: Generate CA private key and self-signed certificate.
+openssl genpkey -algorithm RSA -out $CERT_DIR/ca.key
+openssl req -key $CERT_DIR/ca.key -new -x509 -out $CERT_DIR/ca.crt -days 365000 \
+    -subj "/C=$COUNTRY/ST=$STATE/L=$LOCALITY/O=$ORGANIZATION/OU=$ORG_UNIT/CN=$COMMON_NAME"
+
+# Step 2: Generate server private key.
+openssl genpkey -algorithm RSA -out $CERT_DIR/server.key
+
+# Step 3: Generate server certificate signing request (CSR).
+openssl req -new -key $CERT_DIR/server.key -out $CERT_DIR/server.csr \
+    -subj "/C=$COUNTRY/ST=$STATE/L=$LOCALITY/O=$ORGANIZATION/OU=$ORG_UNIT/CN=$SERVER_CN" \
+    -addext "subjectAltName=DNS:$SERVER_CN,DNS:localhost"
+
+# Step 4: Sign server CSR with the CA certificate to generate the server certificate.
+openssl x509 -req -in $CERT_DIR/server.csr -CA $CERT_DIR/ca.crt -CAkey $CERT_DIR/ca.key \
+    -CAcreateserial -out $CERT_DIR/server.crt -days 365000 \
+    -extfile <(printf "subjectAltName=DNS:$SERVER_CN,DNS:localhost")
+
+# Step 5: Generate client private key.
+openssl genpkey -algorithm RSA -out $CERT_DIR/client.key
+
+# Step 6: Generate client certificate signing request (CSR).
+openssl req -new -key $CERT_DIR/client.key -out $CERT_DIR/client.csr \
+    -subj "/C=$COUNTRY/ST=$STATE/L=$LOCALITY/O=$ORGANIZATION/OU=$ORG_UNIT/CN=$CLIENT_CN"
+
+# Step 7: Sign client CSR with the CA certificate to generate the client certificate.
+openssl x509 -req -in $CERT_DIR/client.csr -CA $CERT_DIR/ca.crt -CAkey $CERT_DIR/ca.key \
+    -CAcreateserial -out $CERT_DIR/client.crt -days 365000
+
+# Step 8: Output summary.
+echo "Certificate Authority (CA) certificate        : $CERT_DIR/ca.crt"
+echo "Server certificate    : $CERT_DIR/server.crt"
+echo "Server private key    : $CERT_DIR/server.key"
+echo "Client certificate    : $CERT_DIR/client.crt"
+echo "Client private key    : $CERT_DIR/client.key"
+
+# Step 9: Remove unused files.
+rm -f $CERT_DIR/*.csr $CERT_DIR/ca.srl $CERT_DIR/ca.key
+
diff --git a/presto-native-execution/presto_cpp/main/connectors/arrow_flight/tests/data/generate_tls_certs.sh b/presto-native-execution/presto_cpp/main/connectors/arrow_flight/tests/data/generate_tls_certs.sh
deleted file mode 100755
index 718f313c70a75..0000000000000
--- a/presto-native-execution/presto_cpp/main/connectors/arrow_flight/tests/data/generate_tls_certs.sh
+++ /dev/null
@@ -1,40 +0,0 @@
-#!/bin/bash
-
-# Set directory for certificates and keys.
-CERT_DIR="./tls_certs"
-mkdir -p $CERT_DIR
-
-# Dummy values for the certificates.
-COUNTRY="US"
-STATE="State"
-LOCALITY="City"
-ORGANIZATION="MyOrg"
-ORG_UNIT="MyUnit"
-COMMON_NAME="MyCA"
-SERVER_CN="server.mydomain.com"
-
-# Step 1: Generate CA private key and self-signed certificate.
-openssl genpkey -algorithm RSA -out $CERT_DIR/ca.key
-openssl req -key $CERT_DIR/ca.key -new -x509 -out $CERT_DIR/ca.crt -days 365000 \
-    -subj "/C=$COUNTRY/ST=$STATE/L=$LOCALITY/O=$ORGANIZATION/OU=$ORG_UNIT/CN=$COMMON_NAME"
-
-# Step 2: Generate server private key.
-openssl genpkey -algorithm RSA -out $CERT_DIR/server.key
-
-# Step 3: Generate server certificate signing request (CSR).
-openssl req -new -key $CERT_DIR/server.key -out $CERT_DIR/server.csr \
-    -subj "/C=$COUNTRY/ST=$STATE/L=$LOCALITY/O=$ORGANIZATION/OU=$ORG_UNIT/CN=$SERVER_CN" \
-    -addext "subjectAltName=DNS:$COMMON_NAME,DNS:localhost" \
-
-# Step 4: Sign server CSR with the CA certificate to generate the server certificate.
-openssl x509 -req -in $CERT_DIR/server.csr -CA $CERT_DIR/ca.crt -CAkey $CERT_DIR/ca.key \
-    -CAcreateserial -out $CERT_DIR/server.crt -days 365000 \
-    -extfile <(printf "subjectAltName=DNS:$COMMON_NAME,DNS:localhost")
-
-# Step 5: Output the generated files.
-echo "Certificate Authority (CA) certificate: $CERT_DIR/ca.crt"
-echo "Server certificate: $CERT_DIR/server.crt"
-echo "Server private key: $CERT_DIR/server.key"
-
-# Step 6: Remove unused files.
-rm -rf $CERT_DIR/server.csr $CERT_DIR/ca.srl $CERT_DIR/ca.key
diff --git a/presto-native-execution/presto_cpp/main/connectors/arrow_flight/tests/data/tls_certs/ca.crt b/presto-native-execution/presto_cpp/main/connectors/arrow_flight/tests/data/tls_certs/ca.crt
deleted file mode 100644
index 6740e89c54e17..0000000000000
--- a/presto-native-execution/presto_cpp/main/connectors/arrow_flight/tests/data/tls_certs/ca.crt
+++ /dev/null
@@ -1,22 +0,0 @@
------BEGIN CERTIFICATE-----
-MIIDmzCCAoOgAwIBAgIUf+rP48iL39yGlAfFQTIp5bmM4uQwDQYJKoZIhvcNAQEL
-BQAwXDELMAkGA1UEBhMCVVMxDjAMBgNVBAgMBVN0YXRlMQ0wCwYDVQQHDARDaXR5
-MQ4wDAYDVQQKDAVNeU9yZzEPMA0GA1UECwwGTXlVbml0MQ0wCwYDVQQDDARNeUNB
-MCAXDTI0MTIwMzExMDQxMVoYDzMwMjQwNDA1MTEwNDExWjBcMQswCQYDVQQGEwJV
-UzEOMAwGA1UECAwFU3RhdGUxDTALBgNVBAcMBENpdHkxDjAMBgNVBAoMBU15T3Jn
-MQ8wDQYDVQQLDAZNeVVuaXQxDTALBgNVBAMMBE15Q0EwggEiMA0GCSqGSIb3DQEB
-AQUAA4IBDwAwggEKAoIBAQCliiXIcSmxXAAq2k/XjcZniDgEDCxWKZGiV8JBiJwY
-MMBJtqcVzWfiDpO2u6d1dfGb6utlRW+1dnwupzURCMmZff4bqlPx4ZejRXDrWzKz
-08WSpDVZwC2H5XOllwK36Cn4gvPRe3YWVcdDGHy7GL+zsJENvawJj0BH952MU4bk
-sV52zEkN291bfN9sSYfT1NCJuLPM0Qsf97DeQ+wHXEw+t4XVMF3FQbciQp0y6CnA
-wfFFN14WDiWxukP1I3kuDYYA6h/WJCQMp5rU2NCB9nIQrulYRxFaepMYENLxgAyj
-gFaoRh2Kt2k7XKv6WOa6CmYm2dZERPlbA+oNAHkaHw6lAgMBAAGjUzBRMB0GA1Ud
-DgQWBBSN+3vRlXGjs6c+rN94qgEnkPLl3DAfBgNVHSMEGDAWgBSN+3vRlXGjs6c+
-rN94qgEnkPLl3DAPBgNVHRMBAf8EBTADAQH/MA0GCSqGSIb3DQEBCwUAA4IBAQAb
-L40Oe2b/2xdUSyjqXJceVxaeA291fCpnu1C1JliP0hNI3fu9jjZhXHQoYub/4mod
-8lriEDIcOCCiUfmi404akpqQHuBmOHaKEOtaaQkezjPsYnUra+O2ssqUo2zto5bK
-gR0LGsb+4AO0bDvq+QVI6kEQqAAIf6qC+kpg/jV4iKJ1J6Qw4R3QppYBm6SQcfvI
-hfUfDSO6SNfy0f/ZVCavbJIP9zG/BfAD9DEERocw03PiN5bm4IXJ3HH8rxyuBfJ5
-Eg/fPP5TlZ2H7Kqb3VgVBGWJtNXWmJphHyraBJTEuxgXWvl6AaW0P/3dsJi3rfdD
-zDIT7AmENLCom8Gl0bgM
------END CERTIFICATE-----
diff --git a/presto-native-execution/presto_cpp/main/connectors/arrow_flight/tests/data/tls_certs/server.crt b/presto-native-execution/presto_cpp/main/connectors/arrow_flight/tests/data/tls_certs/server.crt
deleted file mode 100644
index 92c91f2d613b0..0000000000000
--- a/presto-native-execution/presto_cpp/main/connectors/arrow_flight/tests/data/tls_certs/server.crt
+++ /dev/null
@@ -1,22 +0,0 @@
------BEGIN CERTIFICATE-----
-MIIDtTCCAp2gAwIBAgIUUhmhZP94nIowrg2EarzfEBp6W1EwDQYJKoZIhvcNAQEL
-BQAwXDELMAkGA1UEBhMCVVMxDjAMBgNVBAgMBVN0YXRlMQ0wCwYDVQQHDARDaXR5
-MQ4wDAYDVQQKDAVNeU9yZzEPMA0GA1UECwwGTXlVbml0MQ0wCwYDVQQDDARNeUNB
-MCAXDTI0MTIwMzExMDQxMVoYDzMwMjQwNDA1MTEwNDExWjBrMQswCQYDVQQGEwJV
-UzEOMAwGA1UECAwFU3RhdGUxDTALBgNVBAcMBENpdHkxDjAMBgNVBAoMBU15T3Jn
-MQ8wDQYDVQQLDAZNeVVuaXQxHDAaBgNVBAMME3NlcnZlci5teWRvbWFpbi5jb20w
-ggEiMA0GCSqGSIb3DQEBAQUAA4IBDwAwggEKAoIBAQDSxC4zCC4GFZbX+fdFgWbL
-sj4PortyOM7mzRjNaQ3M0FTSEy5xET9C2qFlBCJ7AL7DlbSLmSckYY/FkdfMqNN2
-+NZ0Dy2d6bZN+ly5N/QBVnyS/5HVC3MXa6Y2BmFXiBnczWfGBwj+uVHlKOUWUyNi
-EyUkhuPwtYXkFmJoqBxJSPC6cxX6NzMujnwCF18dUf0Vra44osu4moaovmg3c9jM
-cBtmafFs9F54FoAEuLotjISVEa7VY6th5RxXJHpgas+0R5EBddGYKbTRiUYjht7r
-pS+An0ey02oOjEWdqLnQSg/SUGKuRXULyE5l1A0HfNQtvepUQotb9ull1F7OrbfB
-AgMBAAGjXjBcMBoGA1UdEQQTMBGCBE15Q0GCCWxvY2FsaG9zdDAdBgNVHQ4EFgQU
-vnCLWjre4jqkKzC24psCPh1oIQwwHwYDVR0jBBgwFoAUjft70ZVxo7OnPqzfeKoB
-J5Dy5dwwDQYJKoZIhvcNAQELBQADggEBAJCiJgtTw/7b/g3QnJfM4teQkFS440Ii
-weqQJMoP6als8Fc3opPKv9eC5w0wqaLlIdwJjzGM5PmCAtGVafo22TbqhZyQdzQu
-TUKv1DaVF0JBVAGVxTSDIK9r5Ww4mDAQnQENLC6soS3AvYDEi+8667YLoNNdhRCX
-q2D5v76UN45idiShppxOw53whsvpHv+wyqcdse7DhgM9boCbx51Uvv3l/AEToyaj
-S1xeIkBwNpSYU0ax2Lr1j2yoKbzAa3MHy8Php+T5CGji02+HwwlvlPDLtw8q5gHw
-BLSwlAHgclPxUTWNNoCqjfX8Bi083+QDCLm0rgQ45xljNDbFAF1Y5hA=
------END CERTIFICATE-----
diff --git a/presto-native-execution/presto_cpp/main/connectors/arrow_flight/tests/data/tls_certs/server.key b/presto-native-execution/presto_cpp/main/connectors/arrow_flight/tests/data/tls_certs/server.key
deleted file mode 100644
index 2cdf5750a4753..0000000000000
--- a/presto-native-execution/presto_cpp/main/connectors/arrow_flight/tests/data/tls_certs/server.key
+++ /dev/null
@@ -1,28 +0,0 @@
------BEGIN PRIVATE KEY-----
-MIIEvAIBADANBgkqhkiG9w0BAQEFAASCBKYwggSiAgEAAoIBAQDSxC4zCC4GFZbX
-+fdFgWbLsj4PortyOM7mzRjNaQ3M0FTSEy5xET9C2qFlBCJ7AL7DlbSLmSckYY/F
-kdfMqNN2+NZ0Dy2d6bZN+ly5N/QBVnyS/5HVC3MXa6Y2BmFXiBnczWfGBwj+uVHl
-KOUWUyNiEyUkhuPwtYXkFmJoqBxJSPC6cxX6NzMujnwCF18dUf0Vra44osu4moao
-vmg3c9jMcBtmafFs9F54FoAEuLotjISVEa7VY6th5RxXJHpgas+0R5EBddGYKbTR
-iUYjht7rpS+An0ey02oOjEWdqLnQSg/SUGKuRXULyE5l1A0HfNQtvepUQotb9ull
-1F7OrbfBAgMBAAECggEAAxbZuuESGGAMMm9HLGhKHgbHU8gnv2Phdbrka+SYBYg5
-UYzTHLh3FwEsjd4VnaweJ7CN1WDb1NvWmTum/DCebJ1HKqtjKLAZfk8q2TLGmXdL
-pzWOdQ8MX1fKP2sIlcl0kFbNCE8vprjneDyBLtqOK36eiAh/fl6BQ12QAMLjyv/L
-OwXSY4ESs/RzxRzFgdT98cDZFL7y0FVIjJo/Q5lfW9UwwSfw8tOLNXKTYwPHqIfJ
-NjfWD7IqztQlnanyRXv5dScp80i8p9qgH0i8YfVBHZDeOmHGLcltilLRZ0dQ/X0g
-Lrr0aIO3iLhmTIkJRzUnGeyvDjxcPINvRSBBwXy04QKBgQDpFJa/EwSsWj8586oh
-xgm0Z3q+FiEeCe7aLLPcXAS2EDvix5ibJDT2y1Aa/kXq25S53npa/7Ov6TJs5H4g
-eyshDtR1wVhz+rIggREiX/sagkhwnNsssUZFv5t9PdnaFXpVnH49m5Qc8HO3owtN
-t8EGSRcAQ4o/fLWLs51qd38cIQKBgQDnfd8YPyDQ03xDC/3+Qrypyc/xhGnCuj7w
-ZeA5iEyTnnNxL0a0B6PWcSk2BZReMNQKgYtipnsOQKtwHMttxtXYs/VQpeB4KoWE
-zEwW0fV3MMsXN+nVJlEZnVaTbmYXknjeZrh/rNjsY96yxw8NtvAuYSpnqtr3N2nd
-iMQ3G/QnoQKBgGMi+bdNvIgeXpQkmrGAzTHpbaCaQv3G1cwAhYPts6dIomAj6znZ
-nZl3ApxomI57VPf1s+8uoVvqASOl0Cu6l66Y4y8uzJOQBuGiZApN7rzouy0C2opY
-4H3cMKOFgjqrNfxh8qP7n3TrpRxvgehNhxFIVzsqfwvf3EwOWp8lMnBhAoGAZ25E
-Ge9K2ENGCCb5i3uCFFLJiF3ja1AQAxVhxBL0NBjd97pp2tJ3D79r7Gk9y4ABndgX
-0TIVVV7ruqIC+r+WmMZ/W1NiIg7NrXIipSeWh3TTqUIgRk5iehFkt2biUrHtM2Gu
-Gc2+9pAA1tw+C6CrW+2qJrueLksiEAulsAHba0ECgYBIgIiY+Gx+XecEgCwAhWcn
-GzNDAAlA4IgBjHpUtIByflzQDqlECKXjPbVBKfyq6eLt40upFmQCLsn+AkiQau8A
-3cFAK9wJOAHv9KuWDrbHyhRE9CrJ6BqsY2goC3LiFCTgJy1TrRl6CDaFzHivONwF
-LNPflYk5s376UWqxC+HtIA==
------END PRIVATE KEY-----
diff --git a/presto-native-execution/presto_cpp/main/operators/BroadcastFactory.cpp b/presto-native-execution/presto_cpp/main/operators/BroadcastFactory.cpp
index 7dcfee2cbfad3..b82ea7192e6b1 100644
--- a/presto-native-execution/presto_cpp/main/operators/BroadcastFactory.cpp
+++ b/presto-native-execution/presto_cpp/main/operators/BroadcastFactory.cpp
@@ -82,7 +82,12 @@ void BroadcastFileWriter::collect(const RowVectorPtr& input) {
   write(input);
 }
 
-void BroadcastFileWriter::noMoreData() {}
+void BroadcastFileWriter::noMoreData() {
+  if (writeFile_ != nullptr) {
+    writeFile_->flush();
+    writeFile_->close();
+  }
+}
 
 RowVectorPtr BroadcastFileWriter::fileStats() {
   // No rows written.
diff --git a/presto-native-execution/presto_cpp/main/operators/BroadcastWrite.cpp b/presto-native-execution/presto_cpp/main/operators/BroadcastWrite.cpp
index cffcbd4a64320..fc835bbe5c73b 100644
--- a/presto-native-execution/presto_cpp/main/operators/BroadcastWrite.cpp
+++ b/presto-native-execution/presto_cpp/main/operators/BroadcastWrite.cpp
@@ -70,6 +70,9 @@ class BroadcastWriteOperator : public Operator {
     }
 
     fileBroadcastWriter_->collect(reorderedInput);
+    auto lockedStats = stats_.wlock();
+    lockedStats->addOutputVector(
+        reorderedInput->estimateFlatSize(), reorderedInput->size());
   }
 
   void noMoreInput() override {
diff --git a/presto-native-execution/presto_cpp/main/operators/ShuffleWrite.cpp b/presto-native-execution/presto_cpp/main/operators/ShuffleWrite.cpp
index 434353a1e5d2e..dbb99c0e62711 100644
--- a/presto-native-execution/presto_cpp/main/operators/ShuffleWrite.cpp
+++ b/presto-native-execution/presto_cpp/main/operators/ShuffleWrite.cpp
@@ -104,6 +104,8 @@ class ShuffleWriteOperator : public Operator {
             "collect");
       }
     }
+    auto lockedStats = stats_.wlock();
+    lockedStats->addOutputVector(input->estimateFlatSize(), input->size());
   }
 
   void noMoreInput() override {
diff --git a/presto-native-execution/presto_cpp/main/operators/tests/BinarySortableSerializerTest.cpp b/presto-native-execution/presto_cpp/main/operators/tests/BinarySortableSerializerTest.cpp
index 4db5793e31708..77481b697befb 100644
--- a/presto-native-execution/presto_cpp/main/operators/tests/BinarySortableSerializerTest.cpp
+++ b/presto-native-execution/presto_cpp/main/operators/tests/BinarySortableSerializerTest.cpp
@@ -1045,7 +1045,7 @@ TEST_F(BinarySortableSerializerTest, ArrayTypeSingleFieldTests) {
   // null < []
   EXPECT_TRUE(
       singleArrayFieldCompare(
-          {std::nullopt, {{}}}, velox::core::kAscNullsFirst) < 0);
+          {std::nullopt, {std::vector>{}}}, velox::core::kAscNullsFirst) < 0);
 }
 
 TEST_F(BinarySortableSerializerTest, RowTypeSingleFieldTests) {
diff --git a/presto-native-execution/presto_cpp/main/operators/tests/BroadcastTest.cpp b/presto-native-execution/presto_cpp/main/operators/tests/BroadcastTest.cpp
index c685ede460fee..95e63ca080251 100644
--- a/presto-native-execution/presto_cpp/main/operators/tests/BroadcastTest.cpp
+++ b/presto-native-execution/presto_cpp/main/operators/tests/BroadcastTest.cpp
@@ -272,7 +272,7 @@ TEST_F(BroadcastTest, endToEndSerdeLayout) {
   runBroadcastTest({data}, {{"c1", "c1", "c2"}});
 
   // Skip all.
-  runBroadcastTest({data}, {{}});
+  runBroadcastTest({data}, {std::vector{}});
 }
 
 TEST_F(BroadcastTest, endToEndWithNoRows) {
diff --git a/presto-native-execution/presto_cpp/main/runtime-metrics/PrometheusStatsReporter.h b/presto-native-execution/presto_cpp/main/runtime-metrics/PrometheusStatsReporter.h
index bf8574a624f4e..c5764d12581f6 100644
--- a/presto-native-execution/presto_cpp/main/runtime-metrics/PrometheusStatsReporter.h
+++ b/presto-native-execution/presto_cpp/main/runtime-metrics/PrometheusStatsReporter.h
@@ -65,6 +65,30 @@ class PrometheusStatsReporter : public facebook::velox::BaseStatsReporter {
       int64_t max,
       const std::vector& pcts) const override;
 
+  void registerQuantileMetricExportType(
+      const char* /* key */,
+      const std::vector& /* statTypes */,
+      const std::vector& /* pcts */,
+      const std::vector& /* slidingWindowsSeconds */) const override {};
+
+  void registerQuantileMetricExportType(
+      folly::StringPiece /* key */,
+      const std::vector& /* statTypes */,
+      const std::vector& /* pcts */,
+      const std::vector& /* slidingWindowsSeconds */) const override {};
+  
+  void registerDynamicQuantileMetricExportType(
+      const char* /* keyPattern */,
+      const std::vector& /* statTypes */,
+      const std::vector& /* pcts */,
+      const std::vector& /* slidingWindowsSeconds */) const override {};
+  
+  void registerDynamicQuantileMetricExportType(
+      folly::StringPiece /* keyPattern */,
+      const std::vector& /* statTypes */,
+      const std::vector& /* pcts */,
+      const std::vector& /* slidingWindowsSeconds */) const override {};
+
   void addMetricValue(const std::string& key, size_t value = 1) const override;
 
   void addMetricValue(const char* key, size_t value = 1) const override;
@@ -79,6 +103,30 @@ class PrometheusStatsReporter : public facebook::velox::BaseStatsReporter {
   void addHistogramMetricValue(folly::StringPiece key, size_t value)
       const override;
 
+  void addQuantileMetricValue(const std::string& /* key */, size_t /* value */)
+      const override {};
+
+  void addQuantileMetricValue(const char* /* key */, size_t /* value */)
+      const override {};
+
+  void addQuantileMetricValue(folly::StringPiece /* key */, size_t /* value */)
+      const override{};
+  
+  void addDynamicQuantileMetricValue(
+      const std::string& /* key */,
+      folly::Range /* subkeys */,
+      size_t /* value */) const override {};
+  
+  virtual void addDynamicQuantileMetricValue(
+      const char* /* key */,
+      folly::Range /* subkeys */,
+      size_t /* value */) const override {};
+
+  virtual void addDynamicQuantileMetricValue(
+      folly::StringPiece /* key */,
+      folly::Range /* subkeys */,
+      size_t /* value */) const override {};
+
   std::string fetchMetrics() override;
 
   /**
diff --git a/presto-native-execution/presto_cpp/main/tests/CMakeLists.txt b/presto-native-execution/presto_cpp/main/tests/CMakeLists.txt
index f95483eb3ad25..2e24268aa1a07 100644
--- a/presto-native-execution/presto_cpp/main/tests/CMakeLists.txt
+++ b/presto-native-execution/presto_cpp/main/tests/CMakeLists.txt
@@ -14,7 +14,6 @@ add_executable(
   AnnouncerTest.cpp
   CoordinatorDiscovererTest.cpp
   HttpServerWrapper.cpp
-  MutableConfigs.cpp
   PeriodicMemoryCheckerTest.cpp
   PrestoExchangeSourceTest.cpp
   PrestoTaskTest.cpp
@@ -52,6 +51,7 @@ target_link_libraries(
   velox_functions_prestosql
   velox_aggregates
   velox_hive_partition_function
+  presto_mutable_configs
   ${RE2}
   GTest::gmock
   GTest::gtest
diff --git a/presto-native-execution/presto_cpp/main/tests/CoordinatorDiscovererTest.cpp b/presto-native-execution/presto_cpp/main/tests/CoordinatorDiscovererTest.cpp
index 76315aeb7bc4a..c942515aa535a 100644
--- a/presto-native-execution/presto_cpp/main/tests/CoordinatorDiscovererTest.cpp
+++ b/presto-native-execution/presto_cpp/main/tests/CoordinatorDiscovererTest.cpp
@@ -15,7 +15,7 @@
 #include "presto_cpp/main/CoordinatorDiscoverer.h"
 #include 
 #include "presto_cpp/main/common/Configs.h"
-#include "presto_cpp/main/tests/MultableConfigs.h"
+#include "presto_cpp/main/common/tests/MutableConfigs.h"
 #include "velox/common/file/FileSystems.h"
 
 using namespace facebook::velox;
diff --git a/presto-native-execution/presto_cpp/main/tests/PrestoExchangeSourceTest.cpp b/presto-native-execution/presto_cpp/main/tests/PrestoExchangeSourceTest.cpp
index 17fe5659484e4..a1e76624a6ebc 100644
--- a/presto-native-execution/presto_cpp/main/tests/PrestoExchangeSourceTest.cpp
+++ b/presto-native-execution/presto_cpp/main/tests/PrestoExchangeSourceTest.cpp
@@ -19,9 +19,9 @@
 #include 
 #include "folly/experimental/EventCount.h"
 #include "presto_cpp/main/PrestoExchangeSource.h"
+#include "presto_cpp/main/common/tests/MutableConfigs.h"
 #include "presto_cpp/main/common/Utils.h"
 #include "presto_cpp/main/tests/HttpServerWrapper.h"
-#include "presto_cpp/main/tests/MultableConfigs.h"
 #include "presto_cpp/presto_protocol/core/presto_protocol_core.h"
 #include "velox/common/base/tests/GTestUtils.h"
 #include "velox/common/file/FileSystems.h"
diff --git a/presto-native-execution/presto_cpp/main/tests/SessionPropertiesTest.cpp b/presto-native-execution/presto_cpp/main/tests/SessionPropertiesTest.cpp
index 20cc3644270f2..c88abee0e3719 100644
--- a/presto-native-execution/presto_cpp/main/tests/SessionPropertiesTest.cpp
+++ b/presto-native-execution/presto_cpp/main/tests/SessionPropertiesTest.cpp
@@ -69,6 +69,8 @@ TEST_F(SessionPropertiesTest, validateMapping) {
        core::QueryConfig::kDebugDisableExpressionWithLazyInputs},
       {SessionProperties::kDebugMemoryPoolNameRegex,
        core::QueryConfig::kDebugMemoryPoolNameRegex},
+      {SessionProperties::kDebugMemoryPoolWarnThresholdBytes,
+       core::QueryConfig::kDebugMemoryPoolWarnThresholdBytes},
       {SessionProperties::kSelectiveNimbleReaderEnabled,
        core::QueryConfig::kSelectiveNimbleReaderEnabled},
       {SessionProperties::kQueryTraceEnabled,
diff --git a/presto-native-execution/presto_cpp/main/tests/TaskManagerTest.cpp b/presto-native-execution/presto_cpp/main/tests/TaskManagerTest.cpp
index bab2407b56d8c..02f84d83a77ed 100644
--- a/presto-native-execution/presto_cpp/main/tests/TaskManagerTest.cpp
+++ b/presto-native-execution/presto_cpp/main/tests/TaskManagerTest.cpp
@@ -18,9 +18,9 @@
 #include "folly/experimental/EventCount.h"
 #include "presto_cpp/main/PrestoExchangeSource.h"
 #include "presto_cpp/main/TaskResource.h"
+#include "presto_cpp/main/common/tests/MutableConfigs.h"
 #include "presto_cpp/main/connectors/PrestoToVeloxConnector.h"
 #include "presto_cpp/main/tests/HttpServerWrapper.h"
-#include "presto_cpp/main/tests/MultableConfigs.h"
 #include "velox/common/base/Fs.h"
 #include "velox/common/base/tests/GTestUtils.h"
 #include "velox/common/file/FileSystems.h"
@@ -715,6 +715,39 @@ TEST_P(TaskManagerTest, tableScanAllSplitsAtOnce) {
   assertResults(taskId, rowType_, "SELECT * FROM tmp WHERE c0 % 5 = 0");
 }
 
+TEST_P(TaskManagerTest, addSplitsWithSameSourceNode) {
+  const auto tableDir = exec::test::TempDirectoryPath::create();
+  auto filePaths = makeFilePaths(tableDir, 5);
+  auto vectors = makeVectors(filePaths.size(), 1'000);
+  for (int i = 0; i < filePaths.size(); i++) {
+    writeToFile(filePaths[i], vectors[i]);
+  }
+  duckDbQueryRunner_.createTable("tmp", vectors);
+
+  const auto planFragment = exec::test::PlanBuilder()
+                          .tableScan(rowType_)
+                          .filter("c0 % 5 = 0")
+                          .partitionedOutput({}, 1, {"c0", "c1"}, GetParam())
+                          .planFragment();
+
+  protocol::TaskUpdateRequest updateRequest;
+  // Create multiple task sources with the same source node id.
+  std::vector taskSources;
+  taskSources.reserve(filePaths.size());
+  long splitSequenceId{0};
+  for (const auto& filePath : filePaths) {
+    taskSources.push_back(makeSource("0", {filePath}, /*noMoreSplits=*/true, splitSequenceId));
+  }
+  taskSources.reserve(filePaths.size());
+  updateRequest.sources = std::move(taskSources);
+
+  protocol::TaskId taskId = "scan.0.0.1.0";
+  auto taskInfo = createOrUpdateTask(taskId, updateRequest, planFragment);
+
+  ASSERT_GE(taskInfo->stats.queuedTimeInNanos, 0);
+  assertResults(taskId, rowType_, "SELECT * FROM tmp WHERE c0 % 5 = 0");
+}
+
 TEST_P(TaskManagerTest, fecthFromFinishedTask) {
   const auto tableDir = exec::test::TempDirectoryPath::create();
   auto filePaths = makeFilePaths(tableDir, 5);
diff --git a/presto-native-execution/presto_cpp/main/types/CMakeLists.txt b/presto-native-execution/presto_cpp/main/types/CMakeLists.txt
index d264bc95f80c2..8b168254334aa 100644
--- a/presto-native-execution/presto_cpp/main/types/CMakeLists.txt
+++ b/presto-native-execution/presto_cpp/main/types/CMakeLists.txt
@@ -11,7 +11,7 @@
 # limitations under the License.
 
 add_library(presto_type_converter OBJECT TypeParser.cpp)
-target_link_libraries(presto_type_converter velox_type_parser)
+target_link_libraries(presto_type_converter velox_presto_type_parser)
 
 add_library(presto_velox_expr_conversion OBJECT PrestoToVeloxExpr.cpp)
 target_link_libraries(presto_velox_expr_conversion velox_presto_types
diff --git a/presto-native-execution/presto_cpp/main/types/PrestoToVeloxQueryPlan.cpp b/presto-native-execution/presto_cpp/main/types/PrestoToVeloxQueryPlan.cpp
index e484dab310148..04dd0b48f3cc6 100644
--- a/presto-native-execution/presto_cpp/main/types/PrestoToVeloxQueryPlan.cpp
+++ b/presto-native-execution/presto_cpp/main/types/PrestoToVeloxQueryPlan.cpp
@@ -419,10 +419,9 @@ core::PlanNodePtr VeloxQueryPlanConverterBase::toVeloxQueryPlan(
 
     const auto desiredSourceOutput = toRowType(node->inputs[i], typeParser_);
 
-    for (auto j = 0; j < outputType->size(); j++) {
-      projections.emplace_back(
-          std::make_shared(
-              outputType->childAt(j), desiredSourceOutput->nameOf(j)));
+    for (auto j = 0; j < outputType->size(); ++j) {
+      projections.emplace_back(std::make_shared(
+          outputType->childAt(j), desiredSourceOutput->nameOf(j)));
     }
 
     sourceNodes[i] = std::make_shared(
@@ -574,6 +573,13 @@ std::shared_ptr isAnd(
   return isSpecialForm(expression, protocol::Form::AND);
 }
 
+// Check if input RowExpression is an 'or' expression and returns it as
+// SpecialFormExpression. Returns nullptr if input expression is something else.
+std::shared_ptr isOr(
+    const std::shared_ptr& expression) {
+  return isSpecialForm(expression, protocol::Form::OR);
+}
+
 // Checks if input PlanNode represents a local exchange with single source and
 // returns it as ExchangeNode. Returns nullptr if input node is something else.
 std::shared_ptr isLocalSingleSourceExchange(
@@ -662,9 +668,8 @@ core::PlanNodePtr VeloxQueryPlanConverterBase::toVeloxQueryPlan(
     std::vector projections;
     projections.reserve(leftNames.size() + 1);
     for (auto i = 0; i < leftNames.size(); i++) {
-      projections.emplace_back(
-          std::make_shared(
-              leftTypes[i], leftNames[i]));
+      projections.emplace_back(std::make_shared(
+          leftTypes[i], leftNames[i]));
     }
     const bool constantValue =
         joinType.value() == core::JoinType::kLeftSemiFilter;
@@ -818,8 +823,8 @@ VectorPtr VeloxQueryPlanConverterBase::evaluateConstantExpression(
   return result[0];
 }
 
-std::shared_ptr
-VeloxQueryPlanConverterBase::generateAggregationNode(
+std::optional
+VeloxQueryPlanConverterBase::toColumnStatsSpec(
     const std::shared_ptr&
         statisticsAggregation,
     core::AggregationNode::Step step,
@@ -828,7 +833,7 @@ VeloxQueryPlanConverterBase::generateAggregationNode(
     const std::shared_ptr& /*tableWriteInfo*/,
     const protocol::TaskId& /*taskId*/) {
   if (statisticsAggregation == nullptr) {
-    return nullptr;
+    return std::nullopt;
   }
   const auto outputVariables = statisticsAggregation->outputVariables;
   const auto aggregationMap = statisticsAggregation->aggregations;
@@ -844,8 +849,7 @@ VeloxQueryPlanConverterBase::generateAggregationNode(
   std::vector aggregateNames;
   std::vector aggregates;
   toAggregations(outputVariables, aggregationMap, aggregates, aggregateNames);
-
-  return std::make_shared(
+  const auto aggregationNode = std::make_shared(
       id,
       step,
       toVeloxExprs(statisticsAggregation->groupingVariables),
@@ -854,6 +858,21 @@ VeloxQueryPlanConverterBase::generateAggregationNode(
       aggregates,
       /*ignoreNullKeys=*/false,
       sourceVeloxPlan);
+
+  // Sanity checks on aggregation node.
+  VELOX_CHECK(!aggregationNode->ignoreNullKeys());
+  VELOX_CHECK(!aggregationNode->groupId().has_value());
+  VELOX_CHECK(!aggregationNode->isPreGrouped());
+  VELOX_CHECK(aggregationNode->globalGroupingSets().empty());
+  VELOX_CHECK(!aggregationNode->aggregateNames().empty());
+  VELOX_CHECK_EQ(
+      aggregationNode->aggregateNames().size(),
+      aggregationNode->aggregates().size());
+  return core::ColumnStatsSpec{
+      aggregationNode->groupingKeys(),
+      aggregationNode->step(),
+      aggregationNode->aggregateNames(),
+      aggregationNode->aggregates()};
 }
 
 std::vector
@@ -1116,9 +1135,8 @@ VeloxQueryPlanConverterBase::toVeloxQueryPlan(
   std::vector groupingKeys;
   groupingKeys.reserve(node->groupingColumns.size());
   for (const auto& [output, input] : node->groupingColumns) {
-    groupingKeys.emplace_back(
-        core::GroupIdNode::GroupingKeyInfo{
-            output.name, exprConverter_.toVeloxExpr(input)});
+    groupingKeys.emplace_back(core::GroupIdNode::GroupingKeyInfo{
+        output.name, exprConverter_.toVeloxExpr(input)});
   }
 
   return std::make_shared(
@@ -1172,6 +1190,17 @@ core::JoinType toJoinType(protocol::JoinType type) {
 
   VELOX_UNSUPPORTED("Unknown join type");
 }
+
+core::JoinType toJoinType(protocol::SpatialJoinType type) {
+  switch (type) {
+    case protocol::SpatialJoinType::INNER:
+      return core::JoinType::kInner;
+    case protocol::SpatialJoinType::LEFT:
+      return core::JoinType::kLeft;
+  }
+
+  VELOX_UNSUPPORTED("Unknown spatial join type");
+}
 } // namespace
 
 core::PlanNodePtr VeloxQueryPlanConverterBase::toVeloxQueryPlan(
@@ -1250,6 +1279,21 @@ core::PlanNodePtr VeloxQueryPlanConverterBase::toVeloxQueryPlan(
       ROW(std::move(outputNames), std::move(outputTypes)));
 }
 
+core::PlanNodePtr VeloxQueryPlanConverterBase::toVeloxQueryPlan(
+    const std::shared_ptr& node,
+    const std::shared_ptr& tableWriteInfo,
+    const protocol::TaskId& taskId) {
+  auto joinType = toJoinType(node->type);
+
+  return std::make_shared(
+      node->id,
+      joinType,
+      exprConverter_.toVeloxExpr(node->filter),
+      toVeloxQueryPlan(node->left, tableWriteInfo, taskId),
+      toVeloxQueryPlan(node->right, tableWriteInfo, taskId),
+      toRowType(node->outputVariables, typeParser_));
+}
+
 std::shared_ptr
 VeloxQueryPlanConverterBase::toVeloxQueryPlan(
     const std::shared_ptr& node,
@@ -1275,12 +1319,25 @@ VeloxQueryPlanConverterBase::toVeloxQueryPlan(
       toVeloxQueryPlan(node->indexSource, tableWriteInfo, taskId);
 
   std::vector joinConditionPtrs{};
+  std::vector unsupportedConditions{};
   if (node->filter) {
     parseIndexLookupCondition(
         *node->filter,
         exprConverter_,
         /*acceptConstant=*/false,
-        joinConditionPtrs);
+        joinConditionPtrs,
+        unsupportedConditions);
+  }
+
+  // Combine unsupported conditions into a single filter using AND
+  core::TypedExprPtr joinFilter;
+  if (!unsupportedConditions.empty()) {
+    if (unsupportedConditions.size() == 1) {
+      joinFilter = unsupportedConditions[0];
+    } else {
+      joinFilter = std::make_shared(
+          BOOLEAN(), unsupportedConditions, "and");
+    }
   }
 
   return std::make_shared(
@@ -1289,6 +1346,7 @@ VeloxQueryPlanConverterBase::toVeloxQueryPlan(
       leftKeys,
       rightKeys,
       joinConditionPtrs,
+      joinFilter,
       false,
       left,
       std::dynamic_pointer_cast(right),
@@ -1451,19 +1509,18 @@ VeloxQueryPlanConverterBase::toVeloxQueryPlan(
       typeParser_);
   const auto sourceVeloxPlan =
       toVeloxQueryPlan(node->source, tableWriteInfo, taskId);
-  std::shared_ptr aggregationNode =
-      generateAggregationNode(
-          node->statisticsAggregation,
-          core::AggregationNode::Step::kPartial,
-          node->id,
-          sourceVeloxPlan,
-          tableWriteInfo,
-          taskId);
+  std::optional columnStatsSpec = toColumnStatsSpec(
+      node->statisticsAggregation,
+      core::AggregationNode::Step::kPartial,
+      node->id,
+      sourceVeloxPlan,
+      tableWriteInfo,
+      taskId);
   return std::make_shared(
       node->id,
       toRowType(node->columns, typeParser_),
       node->columnNames,
-      std::move(aggregationNode),
+      columnStatsSpec,
       std::move(insertTableHandle),
       node->partitioningScheme != nullptr,
       outputType,
@@ -1521,7 +1578,7 @@ VeloxQueryPlanConverterBase::toVeloxQueryPlan(
       node->id,
       inputColumns,
       inputColumns->names(),
-      /*aggregationNode=*/nullptr,
+      /*columnStatsSpec=*/std::nullopt,
       std::move(insertTableHandle),
       true, // delete only supported on partitioned tables
       outputType,
@@ -1543,17 +1600,16 @@ VeloxQueryPlanConverterBase::toVeloxQueryPlan(
       typeParser_);
   const auto sourceVeloxPlan =
       toVeloxQueryPlan(node->source, tableWriteInfo, taskId);
-  std::shared_ptr aggregationNode =
-      generateAggregationNode(
-          node->statisticsAggregation,
-          core::AggregationNode::Step::kIntermediate,
-          node->id,
-          sourceVeloxPlan,
-          tableWriteInfo,
-          taskId);
+  std::optional columnStatsSpec = toColumnStatsSpec(
+      node->statisticsAggregation,
+      core::AggregationNode::Step::kIntermediate,
+      node->id,
+      sourceVeloxPlan,
+      tableWriteInfo,
+      taskId);
 
   return std::make_shared(
-      node->id, outputType, aggregationNode, sourceVeloxPlan);
+      node->id, outputType, columnStatsSpec, sourceVeloxPlan);
 }
 
 std::shared_ptr
@@ -1830,6 +1886,10 @@ core::PlanNodePtr VeloxQueryPlanConverterBase::toVeloxQueryPlan(
           std::dynamic_pointer_cast(node)) {
     return toVeloxQueryPlan(join, tableWriteInfo, taskId);
   }
+  if (auto spatialJoin =
+          std::dynamic_pointer_cast(node)) {
+    return toVeloxQueryPlan(spatialJoin, tableWriteInfo, taskId);
+  }
   if (auto remoteSource =
           std::dynamic_pointer_cast(node)) {
     return toVeloxQueryPlan(remoteSource, tableWriteInfo, taskId);
@@ -2006,7 +2066,7 @@ core::PlanFragment VeloxQueryPlanConverterBase::toVeloxQueryPlan(
         planFragment.planNode = core::PartitionedOutputNode::single(
             partitionedOutputNodeId,
             outputType,
-            toVeloxSerdeKind((partitioningScheme.encoding)),
+            toVeloxSerdeKind(partitioningScheme.encoding),
             sourceNode);
         return planFragment;
       case protocol::SystemPartitioning::FIXED: {
@@ -2020,7 +2080,7 @@ core::PlanFragment VeloxQueryPlanConverterBase::toVeloxQueryPlan(
               planFragment.planNode = core::PartitionedOutputNode::single(
                   partitionedOutputNodeId,
                   outputType,
-                  toVeloxSerdeKind((partitioningScheme.encoding)),
+                  toVeloxSerdeKind(partitioningScheme.encoding),
                   sourceNode);
               return planFragment;
             }
@@ -2033,7 +2093,7 @@ core::PlanFragment VeloxQueryPlanConverterBase::toVeloxQueryPlan(
                     partitioningScheme.replicateNullsAndAny,
                     std::make_shared(),
                     outputType,
-                    toVeloxSerdeKind((partitioningScheme.encoding)),
+                    toVeloxSerdeKind(partitioningScheme.encoding),
                     sourceNode);
             return planFragment;
           }
@@ -2046,7 +2106,7 @@ core::PlanFragment VeloxQueryPlanConverterBase::toVeloxQueryPlan(
               planFragment.planNode = core::PartitionedOutputNode::single(
                   partitionedOutputNodeId,
                   outputType,
-                  toVeloxSerdeKind((partitioningScheme.encoding)),
+                  toVeloxSerdeKind(partitioningScheme.encoding),
                   sourceNode);
               return planFragment;
             }
@@ -2060,7 +2120,7 @@ core::PlanFragment VeloxQueryPlanConverterBase::toVeloxQueryPlan(
                     std::make_shared(
                         inputType, keyChannels, constValues),
                     outputType,
-                    toVeloxSerdeKind((partitioningScheme.encoding)),
+                    toVeloxSerdeKind(partitioningScheme.encoding),
                     sourceNode);
             return planFragment;
           }
@@ -2069,7 +2129,7 @@ core::PlanFragment VeloxQueryPlanConverterBase::toVeloxQueryPlan(
                 partitionedOutputNodeId,
                 1,
                 outputType,
-                toVeloxSerdeKind((partitioningScheme.encoding)),
+                toVeloxSerdeKind(partitioningScheme.encoding),
                 sourceNode);
             return planFragment;
           }
@@ -2088,7 +2148,7 @@ core::PlanFragment VeloxQueryPlanConverterBase::toVeloxQueryPlan(
         planFragment.planNode = core::PartitionedOutputNode::arbitrary(
             partitionedOutputNodeId,
             std::move(outputType),
-            toVeloxSerdeKind((partitioningScheme.encoding)),
+            toVeloxSerdeKind(partitioningScheme.encoding),
             std::move(sourceNode));
         return planFragment;
       }
@@ -2107,7 +2167,7 @@ core::PlanFragment VeloxQueryPlanConverterBase::toVeloxQueryPlan(
     planFragment.planNode = core::PartitionedOutputNode::single(
         partitionedOutputNodeId,
         outputType,
-        toVeloxSerdeKind((partitioningScheme.encoding)),
+        toVeloxSerdeKind(partitioningScheme.encoding),
         sourceNode);
     return planFragment;
   }
@@ -2123,7 +2183,7 @@ core::PlanFragment VeloxQueryPlanConverterBase::toVeloxQueryPlan(
       partitioningScheme.replicateNullsAndAny,
       std::shared_ptr(std::move(spec)),
       toRowType(partitioningScheme.outputLayout, typeParser_),
-      toVeloxSerdeKind((partitioningScheme.encoding)),
+      toVeloxSerdeKind(partitioningScheme.encoding),
       sourceNode);
   return planFragment;
 }
@@ -2302,16 +2362,49 @@ void parseSqlFunctionHandle(
   }
 }
 
+void handleUnsupportedIndexLookupCondition(
+    const std::shared_ptr& filter,
+    const VeloxExprConverter& exprConverter,
+    std::vector& unsupportedConditions) {
+  // OR conditions cannot be converted to index lookup conditions
+  if (const auto orForm = isOr(filter)) {
+    VELOX_UNSUPPORTED(
+        "Unsupported index lookup condition: {}", toJsonString(filter));
+  }
+  unsupportedConditions.push_back(exprConverter.toVeloxExpr(filter));
+}
+
+#ifdef VELOX_ENABLE_BACKWARD_COMPATIBILITY
 void parseIndexLookupCondition(
     const std::shared_ptr& filter,
     const VeloxExprConverter& exprConverter,
     bool acceptConstant,
     std::vector& joinConditionPtrs) {
+  std::vector unsupportedConditions{};
+  parseIndexLookupCondition(
+      filter,
+      exprConverter,
+      acceptConstant,
+      joinConditionPtrs,
+      unsupportedConditions);
+}
+#endif
+
+void parseIndexLookupCondition(
+    const std::shared_ptr& filter,
+    const VeloxExprConverter& exprConverter,
+    bool acceptConstant,
+    std::vector& joinConditionPtrs,
+    std::vector& unsupportedConditions) {
   if (const auto andForm = isAnd(filter)) {
     VELOX_CHECK_EQ(andForm->arguments.size(), 2);
     for (const auto& child : andForm->arguments) {
       parseIndexLookupCondition(
-          child, exprConverter, acceptConstant, joinConditionPtrs);
+          child,
+          exprConverter,
+          acceptConstant,
+          joinConditionPtrs,
+          unsupportedConditions);
     }
     return;
   }
@@ -2327,17 +2420,14 @@ void parseIndexLookupCondition(
     const auto lowerExpr = exprConverter.toVeloxExpr(between->arguments[1]);
     const auto upperExpr = exprConverter.toVeloxExpr(between->arguments[2]);
 
-    VELOX_CHECK(
-        acceptConstant ||
-            !(core::TypedExprs::isConstant(lowerExpr) &&
-              core::TypedExprs::isConstant(upperExpr)),
-        "At least one of the between condition bounds needs to be not constant: {}",
-        toJsonString(filter));
-
-    joinConditionPtrs.push_back(
-        std::make_shared(
-            keyColumnExpr, lowerExpr, upperExpr));
-    return;
+    if (acceptConstant ||
+        !(core::TypedExprs::isConstant(lowerExpr) &&
+          core::TypedExprs::isConstant(upperExpr))) {
+      joinConditionPtrs.push_back(
+          std::make_shared(
+              keyColumnExpr, lowerExpr, upperExpr));
+      return;
+    }
   }
 
   if (const auto contains = isContains(filter)) {
@@ -2350,15 +2440,13 @@ void parseIndexLookupCondition(
 
     const auto conditionColumnExpr =
         exprConverter.toVeloxExpr(contains->arguments[0]);
-    VELOX_CHECK(
-        acceptConstant || !core::TypedExprs::isConstant(conditionColumnExpr),
-        "The condition column needs to be not constant: {}",
-        toJsonString(filter));
 
-    joinConditionPtrs.push_back(
-        std::make_shared(
-            keyColumnExpr, conditionColumnExpr));
-    return;
+    if (acceptConstant || !core::TypedExprs::isConstant(conditionColumnExpr)) {
+      joinConditionPtrs.push_back(
+          std::make_shared(
+              keyColumnExpr, conditionColumnExpr));
+      return;
+    }
   }
 
   if (const auto equals = isEqual(filter)) {
@@ -2369,34 +2457,39 @@ void parseIndexLookupCondition(
     const bool leftIsConstant = core::TypedExprs::isConstant(leftExpr);
     const bool rightIsConstant = core::TypedExprs::isConstant(rightExpr);
 
-    VELOX_CHECK_NE(
-      leftIsConstant,
-      rightIsConstant,
-      "The equal condition must have one key and one constant: {}",
-      toJsonString(filter));
-
-    // Determine which argument is the key (non-constant) and which is the value
-    // (constant)
-    const auto& keyArgument =
-        leftIsConstant ? equals->arguments[1] : equals->arguments[0];
-    const auto& constantExpr = leftIsConstant ? leftExpr : rightExpr;
-
-    const auto keyColumnExpr = exprConverter.toVeloxExpr(
-        std::dynamic_pointer_cast(
-            keyArgument));
+    VELOX_CHECK(
+        !(leftIsConstant && rightIsConstant),
+        "The equal condition must have at least one side to be non-constant: {}",
+        toJsonString(filter));
 
-    VELOX_CHECK_NOT_NULL(
-        keyColumnExpr,
-        "Key argument must be a variable reference: {}",
-        toJsonString(keyArgument));
+    // Check if the equal condition is a constant express and a field access.
+    std::shared_ptr keyArgument;
+    core::TypedExprPtr constantExpr;
+    if (core::TypedExprs::isFieldAccess(leftExpr) && rightIsConstant) {
+      keyArgument = equals->arguments[0];
+      constantExpr = rightExpr;
+    } else if (core::TypedExprs::isFieldAccess(rightExpr) && leftIsConstant) {
+      keyArgument = equals->arguments[1];
+      constantExpr = leftExpr;
+    }
 
-    joinConditionPtrs.push_back(
-        std::make_shared(
-            keyColumnExpr, constantExpr));
-    return;
+    if (keyArgument != nullptr && constantExpr != nullptr) {
+      const auto keyColumnExpr = exprConverter.toVeloxExpr(
+          checked_pointer_cast(
+              keyArgument));
+      VELOX_CHECK_NOT_NULL(
+          keyColumnExpr,
+          "Key argument must be a variable reference: {}",
+          toJsonString(keyArgument));
+      joinConditionPtrs.push_back(
+          std::make_shared(
+              keyColumnExpr, constantExpr));
+      return;
+    }
   }
 
-  VELOX_UNSUPPORTED(
-      "Unsupported index lookup condition: {}", toJsonString(filter));
+  // For unsupported conditions, add to the vector or throw error.
+  handleUnsupportedIndexLookupCondition(
+      filter, exprConverter, unsupportedConditions);
 }
 } // namespace facebook::presto
diff --git a/presto-native-execution/presto_cpp/main/types/PrestoToVeloxQueryPlan.h b/presto-native-execution/presto_cpp/main/types/PrestoToVeloxQueryPlan.h
index 6a728a0b4c44c..1e1880cc9e345 100644
--- a/presto-native-execution/presto_cpp/main/types/PrestoToVeloxQueryPlan.h
+++ b/presto-native-execution/presto_cpp/main/types/PrestoToVeloxQueryPlan.h
@@ -110,15 +110,20 @@ class VeloxQueryPlanConverterBase {
       const std::shared_ptr& tableWriteInfo,
       const protocol::TaskId& taskId);
 
+  velox::core::PlanNodePtr toVeloxQueryPlan(
+      const std::shared_ptr& node,
+      const std::shared_ptr& tableWriteInfo,
+      const protocol::TaskId& taskId);
+
   std::shared_ptr toVeloxQueryPlan(
       const std::shared_ptr& node,
       const std::shared_ptr& tableWriteInfo,
       const protocol::TaskId& taskId);
 
   std::shared_ptr toVeloxQueryPlan(
-    const std::shared_ptr& node,
-    const std::shared_ptr& tableWriteInfo,
-    const protocol::TaskId& taskId);
+      const std::shared_ptr& node,
+      const std::shared_ptr& tableWriteInfo,
+      const protocol::TaskId& taskId);
 
   velox::core::PlanNodePtr toVeloxQueryPlan(
       const std::shared_ptr& node,
@@ -151,9 +156,9 @@ class VeloxQueryPlanConverterBase {
       const protocol::TaskId& taskId);
 
   std::shared_ptr toVeloxQueryPlan(
-    const std::shared_ptr& node,
-    const std::shared_ptr& tableWriteInfo,
-    const protocol::TaskId& taskId);
+      const std::shared_ptr& node,
+      const std::shared_ptr& tableWriteInfo,
+      const protocol::TaskId& taskId);
 
   std::shared_ptr toVeloxQueryPlan(
       const std::shared_ptr& node,
@@ -204,7 +209,7 @@ class VeloxQueryPlanConverterBase {
   velox::VectorPtr evaluateConstantExpression(
       const velox::core::TypedExprPtr& expression);
 
-  std::shared_ptr generateAggregationNode(
+  std::optional toColumnStatsSpec(
       const std::shared_ptr&
           statisticsAggregation,
       velox::core::AggregationNode::Step step,
@@ -292,9 +297,18 @@ void parseSqlFunctionHandle(
     std::vector& rawInputTypes,
     TypeParser& typeParser);
 
+#ifdef VELOX_ENABLE_BACKWARD_COMPATIBILITY
 void parseIndexLookupCondition(
     const std::shared_ptr& filter,
     const VeloxExprConverter& exprConverter,
     bool acceptConstant,
     std::vector& joinConditionPtrs);
+#endif
+
+void parseIndexLookupCondition(
+    const std::shared_ptr& filter,
+    const VeloxExprConverter& exprConverter,
+    bool acceptConstant,
+    std::vector& joinConditionPtrs,
+    std::vector& unsupportedConditions);
 } // namespace facebook::presto
diff --git a/presto-native-execution/presto_cpp/main/types/TypeParser.cpp b/presto-native-execution/presto_cpp/main/types/TypeParser.cpp
index 47e8ce3592255..15d7301c347d4 100644
--- a/presto-native-execution/presto_cpp/main/types/TypeParser.cpp
+++ b/presto-native-execution/presto_cpp/main/types/TypeParser.cpp
@@ -15,17 +15,29 @@
 #include 
 
 #include "presto_cpp/main/types/TypeParser.h"
-#include "velox/type/parser/TypeParser.h"
+#include "velox/functions/prestosql/types/parser/TypeParser.h"
+
+#include "presto_cpp/main/common/Configs.h"
 
 namespace facebook::presto {
 
 velox::TypePtr TypeParser::parse(const std::string& text) const {
+  if (SystemConfig::instance()->charNToVarcharImplicitCast()) {
+    if (text.find("char(") == 0 || text.find("CHAR(") == 0) {
+      return velox::VARCHAR();
+    }
+  }
+  if (!SystemConfig::instance()->enumTypesEnabled()) {
+    if (text.find("BigintEnum") != std::string::npos || text.find("VarcharEnum") != std::string::npos) {
+      VELOX_UNSUPPORTED("Unsupported type: {}", text);
+    }
+  }
   auto it = cache_.find(text);
   if (it != cache_.end()) {
     return it->second;
   }
 
-  auto result = velox::parseType(text);
+  auto result = velox::functions::prestosql::parseType(text);
   cache_.insert({text, result});
   return result;
 }
diff --git a/presto-native-execution/presto_cpp/main/types/tests/CMakeLists.txt b/presto-native-execution/presto_cpp/main/types/tests/CMakeLists.txt
index 27cbf06ffda08..8f98f3f76686f 100644
--- a/presto-native-execution/presto_cpp/main/types/tests/CMakeLists.txt
+++ b/presto-native-execution/presto_cpp/main/types/tests/CMakeLists.txt
@@ -35,7 +35,7 @@ target_link_libraries(
   velox_hive_partition_function)
 
 add_executable(presto_expressions_test RowExpressionTest.cpp ValuesPipeTest.cpp
-                                       PlanConverterTest.cpp)
+                                       PlanConverterTest.cpp TypeParserTest.cpp)
 
 add_test(
   NAME presto_expressions_test
@@ -54,6 +54,7 @@ target_link_libraries(
   $
   $
   presto_operators
+  presto_mutable_configs
   presto_type_test_utils
   velox_core
   velox_dwio_common_exception
@@ -66,9 +67,9 @@ target_link_libraries(
   velox_tpch_connector
   velox_hive_partition_function
   velox_presto_serializer
+  velox_presto_type_parser
   velox_serialization
   velox_type
-  velox_type_parser
   Boost::filesystem
   ${RE2}
   ${FOLLY_WITH_DEPENDENCIES}
diff --git a/presto-native-execution/presto_cpp/main/types/tests/PrestoToVeloxQueryPlanTest.cpp b/presto-native-execution/presto_cpp/main/types/tests/PrestoToVeloxQueryPlanTest.cpp
index 240e9cbd6d75a..3c4045a89b9e9 100644
--- a/presto-native-execution/presto_cpp/main/types/tests/PrestoToVeloxQueryPlanTest.cpp
+++ b/presto-native-execution/presto_cpp/main/types/tests/PrestoToVeloxQueryPlanTest.cpp
@@ -231,6 +231,67 @@ TEST_F(PrestoToVeloxQueryPlanTest, parseIndexJoinNode) {
               "builtInFunctionKind": "ENGINE"
             },
             "returnType": "boolean"
+          },
+          {
+            "@type": "call",
+            "arguments": [
+              {
+                "@type": "constant",
+                "type": "bigint",
+                "valueBlock": "CgAAAExPTkdfQVJSQVkBAAAAAAAAAAAAAAAA"
+              },
+              {
+                "@type": "call",
+                "arguments": [
+                  {
+                    "@type": "variable",
+                    "name": "c1",
+                    "type": "bigint"
+                  },
+                  {
+                    "@type": "constant",
+                    "type": "bigint",
+                    "valueBlock": "CgAAAExPTkdfQVJSQVkBAAAAAGQAAAAAAAAA"
+                  }
+                ],
+                "displayName": "MODULUS",
+                "functionHandle": {
+                  "@type": "$static",
+                  "builtInFunctionKind": "ENGINE",
+                  "signature": {
+                    "argumentTypes": [
+                      "bigint",
+                      "bigint"
+                    ],
+                    "kind": "SCALAR",
+                    "longVariableConstraints": [],
+                    "name": "presto.default.$operator$modulus",
+                    "returnType": "bigint",
+                    "typeVariableConstraints": [],
+                    "variableArity": false
+                  }
+                },
+                "returnType": "bigint"
+              }
+            ],
+            "displayName": "$operator$equal",
+            "functionHandle": {
+              "@type": "$static",
+              "builtInFunctionKind": "ENGINE",
+              "signature": {
+                "argumentTypes": [
+                  "bigint",
+                  "bigint"
+                ],
+                "kind": "SCALAR",
+                "longVariableConstraints": [],
+                "name": "presto.default.$operator$equal",
+                "returnType": "boolean",
+                "typeVariableConstraints": [],
+                "variableArity": false
+              }
+            },
+            "returnType": "boolean"
           }
         ],
         "form": "AND",
diff --git a/presto-native-execution/presto_cpp/main/types/tests/RowExpressionTest.cpp b/presto-native-execution/presto_cpp/main/types/tests/RowExpressionTest.cpp
index 75a28f4db2d5f..6cb1d7d100815 100644
--- a/presto-native-execution/presto_cpp/main/types/tests/RowExpressionTest.cpp
+++ b/presto-native-execution/presto_cpp/main/types/tests/RowExpressionTest.cpp
@@ -14,10 +14,12 @@
 #include 
 #include 
 
+#include "presto_cpp/main/common/Configs.h"
+#include "presto_cpp/main/common/tests/MutableConfigs.h"
 #include "presto_cpp/main/types/PrestoToVeloxExpr.h"
 #include "presto_cpp/presto_protocol/core/presto_protocol_core.h"
+#include "velox/common/file/FileSystems.h"
 #include "velox/core/Expressions.h"
-#include "velox/type/Type.h"
 #include "velox/functions/prestosql/types/JsonRegistration.h"
 
 using namespace facebook::presto;
@@ -32,6 +34,8 @@ class RowExpressionTest : public ::testing::Test {
 
   void SetUp() override {
     registerJsonType();
+    filesystems::registerLocalFileSystem();
+    test::setupMutableSystemConfig();
     pool_ = memory::MemoryManager::getInstance()->addLeafPool();
     converter_ =
         std::make_unique(pool_.get(), &typeParser_);
@@ -414,6 +418,18 @@ TEST_F(RowExpressionTest, varbinary5) {
           '"');
 }
 
+TEST_F(RowExpressionTest, char) {
+  SystemConfig::instance()->setValue(std::string(SystemConfig::kCharNToVarcharImplicitCast), "true");
+  std::string str = R"##(
+        {
+            "@type": "constant",
+            "type": "char(3)",
+            "valueBlock": "DgAAAFZBUklBQkxFX1dJRFRIAQAAAAMAAAAAAwAAAGFiYw=="
+        }
+    )##";
+  testConstantExpression(str, "VARCHAR", "\"abc\"");
+}
+
 TEST_F(RowExpressionTest, timestamp) {
   std::string str = R"(
         {
diff --git a/presto-native-execution/presto_cpp/main/types/tests/TypeParserTest.cpp b/presto-native-execution/presto_cpp/main/types/tests/TypeParserTest.cpp
new file mode 100644
index 0000000000000..ee9036e5943bc
--- /dev/null
+++ b/presto-native-execution/presto_cpp/main/types/tests/TypeParserTest.cpp
@@ -0,0 +1,63 @@
+/*
+ * Licensed 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.
+ */
+
+#include 
+
+#include "presto_cpp/main/common/Configs.h"
+#include "presto_cpp/main/common/tests/MutableConfigs.h"
+#include "presto_cpp/main/types/TypeParser.h"
+#include "velox/common/file/FileSystems.h"
+#include "velox/common/base/tests/GTestUtils.h"
+#include "velox/functions/prestosql/types/BigintEnumRegistration.h"
+#include "velox/functions/prestosql/types/VarcharEnumRegistration.h"
+
+
+using namespace facebook::presto;
+using namespace facebook::velox;
+
+class TypeParserTest : public ::testing::Test {
+  void SetUp() override {
+    filesystems::registerLocalFileSystem();
+    test::setupMutableSystemConfig(); 
+    registerBigintEnumType();
+    registerVarcharEnumType();
+  }
+};
+
+// Test basical functionality of TypeParser.
+// More detailed tests for Presto TypeParser are in velox/functions/prestosql/types/parser/tests/TypeParserTest.
+TEST_F(TypeParserTest, parseEnumTypes) {
+  TypeParser typeParser = TypeParser();
+
+  ASSERT_EQ(
+    typeParser.parse(
+        "test.enum.mood:BigintEnum(test.enum.mood{\"CURIOUS\":2, \"HAPPY\":0})")->toString(),
+    "test.enum.mood:BigintEnum({\"CURIOUS\": 2, \"HAPPY\": 0})");
+  ASSERT_EQ(
+    typeParser.parse(
+        "test.enum.mood:VarcharEnum(test.enum.mood{\"CURIOUS\":\"ONXW2ZKWMFWHKZI=\", \"HAPPY\":\"ONXW2ZJAOZQWY5LF\" , \"SAD\":\"KNHU2RJAKZAUYVKF\"})")->toString(),
+    "test.enum.mood:VarcharEnum({\"CURIOUS\": \"someValue\", \"HAPPY\": \"some value\", \"SAD\": \"SOME VALUE\"})");
+
+  // When set to false, TypeParser will throw an unsupported error when it receives an enum type.
+  SystemConfig::instance()->setValue(std::string(SystemConfig::kEnumTypesEnabled), "false");
+
+  VELOX_ASSERT_THROW(
+    typeParser.parse(
+        "test.enum.mood:BigintEnum(test.enum.mood{\"CURIOUS\":2, \"HAPPY\":0})"),
+    "Unsupported type: test.enum.mood:BigintEnum(test.enum.mood{\"CURIOUS\":2, \"HAPPY\":0})");
+  VELOX_ASSERT_THROW(
+    typeParser.parse(
+        "test.enum.mood:VarcharEnum(test.enum.mood{\"CURIOUS\":\"ONXW2ZKWMFWHKZI=\", \"HAPPY\":\"ONXW2ZJAOZQWY5LF\" , \"SAD\":\"KNHU2RJAKZAUYVKF\"})"),
+    "Unsupported type: test.enum.mood:VarcharEnum(test.enum.mood{\"CURIOUS\":\"ONXW2ZKWMFWHKZI=\", \"HAPPY\":\"ONXW2ZJAOZQWY5LF\" , \"SAD\":\"KNHU2RJAKZAUYVKF\"})");
+}
diff --git a/presto-native-execution/presto_cpp/presto_protocol/connector/iceberg/special/IcebergOutputTableHandle.hpp.inc b/presto-native-execution/presto_cpp/presto_protocol/connector/iceberg/special/IcebergOutputTableHandle.hpp.inc
index 406238e860df3..e168a870cea33 100644
--- a/presto-native-execution/presto_cpp/presto_protocol/connector/iceberg/special/IcebergOutputTableHandle.hpp.inc
+++ b/presto-native-execution/presto_cpp/presto_protocol/connector/iceberg/special/IcebergOutputTableHandle.hpp.inc
@@ -12,7 +12,7 @@
  * limitations under the License.
  */
 
-// IcebergInsertTableHandle is special since it needs an usage of
+// IcebergOutputTableHandle is special since it needs an usage of
 // hive::.
 
 namespace facebook::presto::protocol::iceberg {
diff --git a/presto-native-execution/presto_cpp/presto_protocol/core/Duration.h b/presto-native-execution/presto_cpp/presto_protocol/core/Duration.h
index e0aae14020320..0a5bf10f0589f 100644
--- a/presto-native-execution/presto_cpp/presto_protocol/core/Duration.h
+++ b/presto-native-execution/presto_cpp/presto_protocol/core/Duration.h
@@ -13,6 +13,7 @@
  */
 #pragma once
 #include 
+#include 
 
 namespace facebook::presto::protocol {
 
diff --git a/presto-native-execution/presto_cpp/presto_protocol/core/presto_protocol_core.cpp b/presto-native-execution/presto_cpp/presto_protocol/core/presto_protocol_core.cpp
index fd5029dfe20e1..6f0b26946a186 100644
--- a/presto-native-execution/presto_cpp/presto_protocol/core/presto_protocol_core.cpp
+++ b/presto-native-execution/presto_cpp/presto_protocol/core/presto_protocol_core.cpp
@@ -728,6 +728,10 @@ void to_json(json& j, const std::shared_ptr& p) {
     j = *std::static_pointer_cast(p);
     return;
   }
+  if (type == ".SpatialJoinNode") {
+    j = *std::static_pointer_cast(p);
+    return;
+  }
   if (type == ".TableScanNode") {
     j = *std::static_pointer_cast(p);
     return;
@@ -896,6 +900,12 @@ void from_json(const json& j, std::shared_ptr& p) {
     p = std::static_pointer_cast(k);
     return;
   }
+  if (type == ".SpatialJoinNode") {
+    std::shared_ptr k = std::make_shared();
+    j.get_to(*k);
+    p = std::static_pointer_cast(k);
+    return;
+  }
   if (type == ".TableScanNode") {
     std::shared_ptr k = std::make_shared();
     j.get_to(*k);
@@ -7034,6 +7044,148 @@ void from_json(const json& j, MergeJoinNode& p) {
 } // namespace facebook::presto::protocol
 namespace facebook::presto::protocol {
 
+void to_json(json& j, const NodeLoadMetrics& p) {
+  j = json::object();
+  to_json_key(
+      j,
+      "cpuUsedPercent",
+      p.cpuUsedPercent,
+      "NodeLoadMetrics",
+      "double",
+      "cpuUsedPercent");
+  to_json_key(
+      j,
+      "memoryUsedInBytes",
+      p.memoryUsedInBytes,
+      "NodeLoadMetrics",
+      "double",
+      "memoryUsedInBytes");
+  to_json_key(
+      j,
+      "numQueuedDrivers",
+      p.numQueuedDrivers,
+      "NodeLoadMetrics",
+      "int",
+      "numQueuedDrivers");
+  to_json_key(
+      j,
+      "cpuOverload",
+      p.cpuOverload,
+      "NodeLoadMetrics",
+      "bool",
+      "cpuOverload");
+  to_json_key(
+      j,
+      "memoryOverload",
+      p.memoryOverload,
+      "NodeLoadMetrics",
+      "bool",
+      "memoryOverload");
+}
+
+void from_json(const json& j, NodeLoadMetrics& p) {
+  from_json_key(
+      j,
+      "cpuUsedPercent",
+      p.cpuUsedPercent,
+      "NodeLoadMetrics",
+      "double",
+      "cpuUsedPercent");
+  from_json_key(
+      j,
+      "memoryUsedInBytes",
+      p.memoryUsedInBytes,
+      "NodeLoadMetrics",
+      "double",
+      "memoryUsedInBytes");
+  from_json_key(
+      j,
+      "numQueuedDrivers",
+      p.numQueuedDrivers,
+      "NodeLoadMetrics",
+      "int",
+      "numQueuedDrivers");
+  from_json_key(
+      j,
+      "cpuOverload",
+      p.cpuOverload,
+      "NodeLoadMetrics",
+      "bool",
+      "cpuOverload");
+  from_json_key(
+      j,
+      "memoryOverload",
+      p.memoryOverload,
+      "NodeLoadMetrics",
+      "bool",
+      "memoryOverload");
+}
+} // namespace facebook::presto::protocol
+namespace facebook::presto::protocol {
+// Loosly copied this here from NLOHMANN_JSON_SERIALIZE_ENUM()
+
+// NOLINTNEXTLINE: cppcoreguidelines-avoid-c-arrays
+static const std::pair NodeState_enum_table[] =
+    { // NOLINT: cert-err58-cpp
+        {NodeState::ACTIVE, "ACTIVE"},
+        {NodeState::INACTIVE, "INACTIVE"},
+        {NodeState::SHUTTING_DOWN, "SHUTTING_DOWN"}};
+void to_json(json& j, const NodeState& e) {
+  static_assert(std::is_enum::value, "NodeState must be an enum!");
+  const auto* it = std::find_if(
+      std::begin(NodeState_enum_table),
+      std::end(NodeState_enum_table),
+      [e](const std::pair& ej_pair) -> bool {
+        return ej_pair.first == e;
+      });
+  j = ((it != std::end(NodeState_enum_table))
+           ? it
+           : std::begin(NodeState_enum_table))
+          ->second;
+}
+void from_json(const json& j, NodeState& e) {
+  static_assert(std::is_enum::value, "NodeState must be an enum!");
+  const auto* it = std::find_if(
+      std::begin(NodeState_enum_table),
+      std::end(NodeState_enum_table),
+      [&j](const std::pair& ej_pair) -> bool {
+        return ej_pair.second == j;
+      });
+  e = ((it != std::end(NodeState_enum_table))
+           ? it
+           : std::begin(NodeState_enum_table))
+          ->first;
+}
+} // namespace facebook::presto::protocol
+namespace facebook::presto::protocol {
+
+void to_json(json& j, const NodeStats& p) {
+  j = json::object();
+  to_json_key(
+      j, "nodeState", p.nodeState, "NodeStats", "NodeState", "nodeState");
+  to_json_key(
+      j,
+      "loadMetrics",
+      p.loadMetrics,
+      "NodeStats",
+      "NodeLoadMetrics",
+      "loadMetrics");
+}
+
+void from_json(const json& j, NodeStats& p) {
+  from_json_key(
+      j, "nodeState", p.nodeState, "NodeStats", "NodeState", "nodeState");
+  from_json_key(
+      j,
+      "loadMetrics",
+      p.loadMetrics,
+      "NodeStats",
+      "NodeLoadMetrics",
+      "loadMetrics");
+}
+} // namespace facebook::presto::protocol
+namespace facebook::presto::protocol {
+
 void to_json(json& j, const NodeVersion& p) {
   j = json::object();
   to_json_key(j, "version", p.version, "NodeVersion", "String", "version");
@@ -9343,6 +9495,115 @@ void from_json(const json& j, SortedRangeSet& p) {
 namespace facebook::presto::protocol {
 // Loosly copied this here from NLOHMANN_JSON_SERIALIZE_ENUM()
 
+// NOLINTNEXTLINE: cppcoreguidelines-avoid-c-arrays
+static const std::pair SpatialJoinType_enum_table[] =
+    { // NOLINT: cert-err58-cpp
+        {SpatialJoinType::INNER, "INNER"},
+        {SpatialJoinType::LEFT, "LEFT"}};
+void to_json(json& j, const SpatialJoinType& e) {
+  static_assert(
+      std::is_enum::value, "SpatialJoinType must be an enum!");
+  const auto* it = std::find_if(
+      std::begin(SpatialJoinType_enum_table),
+      std::end(SpatialJoinType_enum_table),
+      [e](const std::pair& ej_pair) -> bool {
+        return ej_pair.first == e;
+      });
+  j = ((it != std::end(SpatialJoinType_enum_table))
+           ? it
+           : std::begin(SpatialJoinType_enum_table))
+          ->second;
+}
+void from_json(const json& j, SpatialJoinType& e) {
+  static_assert(
+      std::is_enum::value, "SpatialJoinType must be an enum!");
+  const auto* it = std::find_if(
+      std::begin(SpatialJoinType_enum_table),
+      std::end(SpatialJoinType_enum_table),
+      [&j](const std::pair& ej_pair) -> bool {
+        return ej_pair.second == j;
+      });
+  e = ((it != std::end(SpatialJoinType_enum_table))
+           ? it
+           : std::begin(SpatialJoinType_enum_table))
+          ->first;
+}
+} // namespace facebook::presto::protocol
+namespace facebook::presto::protocol {
+SpatialJoinNode::SpatialJoinNode() noexcept {
+  _type = ".SpatialJoinNode";
+}
+
+void to_json(json& j, const SpatialJoinNode& p) {
+  j = json::object();
+  j["@type"] = ".SpatialJoinNode";
+  to_json_key(j, "id", p.id, "SpatialJoinNode", "PlanNodeId", "id");
+  to_json_key(j, "type", p.type, "SpatialJoinNode", "SpatialJoinType", "type");
+  to_json_key(j, "left", p.left, "SpatialJoinNode", "PlanNode", "left");
+  to_json_key(j, "right", p.right, "SpatialJoinNode", "PlanNode", "right");
+  to_json_key(
+      j,
+      "outputVariables",
+      p.outputVariables,
+      "SpatialJoinNode",
+      "List",
+      "outputVariables");
+  to_json_key(
+      j, "filter", p.filter, "SpatialJoinNode", "RowExpression", "filter");
+  to_json_key(
+      j,
+      "leftPartitionVariable",
+      p.leftPartitionVariable,
+      "SpatialJoinNode",
+      "VariableReferenceExpression",
+      "leftPartitionVariable");
+  to_json_key(
+      j,
+      "rightPartitionVariable",
+      p.rightPartitionVariable,
+      "SpatialJoinNode",
+      "VariableReferenceExpression",
+      "rightPartitionVariable");
+  to_json_key(j, "kdbTree", p.kdbTree, "SpatialJoinNode", "String", "kdbTree");
+}
+
+void from_json(const json& j, SpatialJoinNode& p) {
+  p._type = j["@type"];
+  from_json_key(j, "id", p.id, "SpatialJoinNode", "PlanNodeId", "id");
+  from_json_key(
+      j, "type", p.type, "SpatialJoinNode", "SpatialJoinType", "type");
+  from_json_key(j, "left", p.left, "SpatialJoinNode", "PlanNode", "left");
+  from_json_key(j, "right", p.right, "SpatialJoinNode", "PlanNode", "right");
+  from_json_key(
+      j,
+      "outputVariables",
+      p.outputVariables,
+      "SpatialJoinNode",
+      "List",
+      "outputVariables");
+  from_json_key(
+      j, "filter", p.filter, "SpatialJoinNode", "RowExpression", "filter");
+  from_json_key(
+      j,
+      "leftPartitionVariable",
+      p.leftPartitionVariable,
+      "SpatialJoinNode",
+      "VariableReferenceExpression",
+      "leftPartitionVariable");
+  from_json_key(
+      j,
+      "rightPartitionVariable",
+      p.rightPartitionVariable,
+      "SpatialJoinNode",
+      "VariableReferenceExpression",
+      "rightPartitionVariable");
+  from_json_key(
+      j, "kdbTree", p.kdbTree, "SpatialJoinNode", "String", "kdbTree");
+}
+} // namespace facebook::presto::protocol
+namespace facebook::presto::protocol {
+// Loosly copied this here from NLOHMANN_JSON_SERIALIZE_ENUM()
+
 // NOLINTNEXTLINE: cppcoreguidelines-avoid-c-arrays
 static const std::pair Form_enum_table[] =
     { // NOLINT: cert-err58-cpp
@@ -11569,39 +11830,3 @@ void from_json(const json& j, WindowNode& p) {
       "preSortedOrderPrefix");
 }
 } // namespace facebook::presto::protocol
-namespace facebook::presto::protocol {
-// Loosly copied this here from NLOHMANN_JSON_SERIALIZE_ENUM()
-
-// NOLINTNEXTLINE: cppcoreguidelines-avoid-c-arrays
-static const std::pair NodeState_enum_table[] =
-    { // NOLINT: cert-err58-cpp
-        {NodeState::ACTIVE, "ACTIVE"},
-        {NodeState::INACTIVE, "INACTIVE"},
-        {NodeState::SHUTTING_DOWN, "SHUTTING_DOWN"}};
-void to_json(json& j, const NodeState& e) {
-  static_assert(std::is_enum::value, "NodeState must be an enum!");
-  const auto* it = std::find_if(
-      std::begin(NodeState_enum_table),
-      std::end(NodeState_enum_table),
-      [e](const std::pair& ej_pair) -> bool {
-        return ej_pair.first == e;
-      });
-  j = ((it != std::end(NodeState_enum_table))
-           ? it
-           : std::begin(NodeState_enum_table))
-          ->second;
-}
-void from_json(const json& j, NodeState& e) {
-  static_assert(std::is_enum::value, "NodeState must be an enum!");
-  const auto* it = std::find_if(
-      std::begin(NodeState_enum_table),
-      std::end(NodeState_enum_table),
-      [&j](const std::pair& ej_pair) -> bool {
-        return ej_pair.second == j;
-      });
-  e = ((it != std::end(NodeState_enum_table))
-           ? it
-           : std::begin(NodeState_enum_table))
-          ->first;
-}
-} // namespace facebook::presto::protocol
diff --git a/presto-native-execution/presto_cpp/presto_protocol/core/presto_protocol_core.h b/presto-native-execution/presto_cpp/presto_protocol/core/presto_protocol_core.h
index 946015a8a6f94..6abc39d461005 100644
--- a/presto-native-execution/presto_cpp/presto_protocol/core/presto_protocol_core.h
+++ b/presto-native-execution/presto_cpp/presto_protocol/core/presto_protocol_core.h
@@ -1761,6 +1761,30 @@ void to_json(json& j, const MergeJoinNode& p);
 void from_json(const json& j, MergeJoinNode& p);
 } // namespace facebook::presto::protocol
 namespace facebook::presto::protocol {
+struct NodeLoadMetrics {
+  double cpuUsedPercent = {};
+  double memoryUsedInBytes = {};
+  int numQueuedDrivers = {};
+  bool cpuOverload = {};
+  bool memoryOverload = {};
+};
+void to_json(json& j, const NodeLoadMetrics& p);
+void from_json(const json& j, NodeLoadMetrics& p);
+} // namespace facebook::presto::protocol
+namespace facebook::presto::protocol {
+enum class NodeState { ACTIVE, INACTIVE, SHUTTING_DOWN };
+extern void to_json(json& j, const NodeState& e);
+extern void from_json(const json& j, NodeState& e);
+} // namespace facebook::presto::protocol
+namespace facebook::presto::protocol {
+struct NodeStats {
+  NodeState nodeState = {};
+  std::shared_ptr loadMetrics = {};
+};
+void to_json(json& j, const NodeStats& p);
+void from_json(const json& j, NodeStats& p);
+} // namespace facebook::presto::protocol
+namespace facebook::presto::protocol {
 struct NodeVersion {
   String version = {};
 };
@@ -2164,6 +2188,27 @@ void to_json(json& j, const SortedRangeSet& p);
 void from_json(const json& j, SortedRangeSet& p);
 } // namespace facebook::presto::protocol
 namespace facebook::presto::protocol {
+enum class SpatialJoinType { INNER, LEFT };
+extern void to_json(json& j, const SpatialJoinType& e);
+extern void from_json(const json& j, SpatialJoinType& e);
+} // namespace facebook::presto::protocol
+namespace facebook::presto::protocol {
+struct SpatialJoinNode : public PlanNode {
+  SpatialJoinType type = {};
+  std::shared_ptr left = {};
+  std::shared_ptr right = {};
+  List outputVariables = {};
+  std::shared_ptr filter = {};
+  std::shared_ptr leftPartitionVariable = {};
+  std::shared_ptr rightPartitionVariable = {};
+  std::shared_ptr kdbTree = {};
+
+  SpatialJoinNode() noexcept;
+};
+void to_json(json& j, const SpatialJoinNode& p);
+void from_json(const json& j, SpatialJoinNode& p);
+} // namespace facebook::presto::protocol
+namespace facebook::presto::protocol {
 enum class Form {
   IF,
   NULL_IF,
@@ -2527,8 +2572,3 @@ struct WindowNode : public PlanNode {
 void to_json(json& j, const WindowNode& p);
 void from_json(const json& j, WindowNode& p);
 } // namespace facebook::presto::protocol
-namespace facebook::presto::protocol {
-enum class NodeState { ACTIVE, INACTIVE, SHUTTING_DOWN };
-extern void to_json(json& j, const NodeState& e);
-extern void from_json(const json& j, NodeState& e);
-} // namespace facebook::presto::protocol
diff --git a/presto-native-execution/presto_cpp/presto_protocol/core/presto_protocol_core.yml b/presto-native-execution/presto_cpp/presto_protocol/core/presto_protocol_core.yml
index 7f90f5d510c9b..aec2731384959 100644
--- a/presto-native-execution/presto_cpp/presto_protocol/core/presto_protocol_core.yml
+++ b/presto-native-execution/presto_cpp/presto_protocol/core/presto_protocol_core.yml
@@ -42,9 +42,6 @@ ExtraFields:
   RemoteTransactionHandle:
     Optional: dummy
 
-AddToOutput:
-  - NodeState
-
 AbstractClasses:
   ColumnHandle:
     super: JsonEncodedSubclass
@@ -160,6 +157,7 @@ AbstractClasses:
       - { name: RemoteSourceNode,         key: com.facebook.presto.sql.planner.plan.RemoteSourceNode }
       - { name: SampleNode,               key: com.facebook.presto.sql.planner.plan.SampleNode }
       - { name: SemiJoinNode,             key: .SemiJoinNode }
+      - { name: SpatialJoinNode,          key: .SpatialJoinNode }
       - { name: TableScanNode,            key: .TableScanNode }
       - { name: TableWriterNode,          key: .TableWriterNode }
       - { name: TableWriterMergeNode,     key: com.facebook.presto.sql.planner.plan.TableWriterMergeNode }
@@ -320,6 +318,7 @@ JavaClasses:
   - presto-spi/src/main/java/com/facebook/presto/spi/plan/JoinNode.java
   - presto-spi/src/main/java/com/facebook/presto/spi/plan/SemiJoinNode.java
   - presto-spi/src/main/java/com/facebook/presto/spi/plan/MergeJoinNode.java
+  - presto-spi/src/main/java/com/facebook/presto/spi/plan/SpatialJoinNode.java
   - presto-main-base/src/main/java/com/facebook/presto/sql/planner/plan/IndexJoinNode.java
   - presto-spi/src/main/java/com/facebook/presto/spi/plan/IndexSourceNode.java
   - presto-spi/src/main/java/com/facebook/presto/spi/plan/TopNNode.java
@@ -345,3 +344,5 @@ JavaClasses:
   - presto-spi/src/main/java/com/facebook/presto/spi/plan/DeleteNode.java
   - presto-spi/src/main/java/com/facebook/presto/spi/plan/BaseInputDistribution.java
   - presto-main-base/src/main/java/com/facebook/presto/metadata/BuiltInFunctionKind.java
+  - presto-spi/src/main/java/com/facebook/presto/spi/NodeStats.java
+  - presto-spi/src/main/java/com/facebook/presto/spi/NodeLoadMetrics.java
diff --git a/presto-native-execution/presto_cpp/presto_protocol/presto_protocol.yml b/presto-native-execution/presto_cpp/presto_protocol/presto_protocol.yml
index 5b407cb8df8b9..219f9bf4bc23a 100644
--- a/presto-native-execution/presto_cpp/presto_protocol/presto_protocol.yml
+++ b/presto-native-execution/presto_cpp/presto_protocol/presto_protocol.yml
@@ -155,6 +155,7 @@ AbstractClasses:
       - { name: RemoteSourceNode,         key: com.facebook.presto.sql.planner.plan.RemoteSourceNode }
       - { name: SampleNode,               key: com.facebook.presto.sql.planner.plan.SampleNode }
       - { name: SemiJoinNode,             key: .SemiJoinNode }
+      - { name: SpatialJoinNode,          key: .SpatialJoinNode }
       - { name: TableScanNode,            key: .TableScanNode }
       - { name: TableWriterNode,          key: .TableWriterNode }
       - { name: TableWriterMergeNode,     key: com.facebook.presto.sql.planner.plan.TableWriterMergeNode }
@@ -360,6 +361,7 @@ JavaClasses:
   - presto-spi/src/main/java/com/facebook/presto/spi/plan/JoinNode.java
   - presto-spi/src/main/java/com/facebook/presto/spi/plan/SemiJoinNode.java
   - presto-spi/src/main/java/com/facebook/presto/spi/plan/MergeJoinNode.java
+  - presto-spi/src/main/java/com/facebook/presto/spi/plan/SpatialJoinNode.java
   - presto-spi/src/main/java/com/facebook/presto/spi/plan/TopNNode.java
   - presto-hive/src/main/java/com/facebook/presto/hive/HivePartitioningHandle.java
   - presto-main/src/main/java/com/facebook/presto/split/EmptySplit.java
diff --git a/presto-native-execution/scripts/dockerfiles/centos-dependency.dockerfile b/presto-native-execution/scripts/dockerfiles/centos-dependency.dockerfile
index e67770e78a0b8..fc9c0c810f265 100644
--- a/presto-native-execution/scripts/dockerfiles/centos-dependency.dockerfile
+++ b/presto-native-execution/scripts/dockerfiles/centos-dependency.dockerfile
@@ -12,9 +12,14 @@
 
 FROM quay.io/centos/centos:stream9
 
-ENV PROMPT_ALWAYS_RESPOND=n
+# Set this when build arm with common flags
+# from https://github.com/facebookincubator/velox/pull/14366
+ARG ARM_BUILD_TARGET
+
+ENV PROMPT_ALWAYS_RESPOND=y
 ENV CC=/opt/rh/gcc-toolset-12/root/bin/gcc
 ENV CXX=/opt/rh/gcc-toolset-12/root/bin/g++
+ENV ARM_BUILD_TARGET=${ARM_BUILD_TARGET}
 
 RUN mkdir -p /scripts /velox/scripts
 COPY scripts /scripts
diff --git a/presto-native-execution/scripts/dockerfiles/prestissimo-runtime.dockerfile b/presto-native-execution/scripts/dockerfiles/prestissimo-runtime.dockerfile
index 30f299172d7fe..826a818b0d7f5 100644
--- a/presto-native-execution/scripts/dockerfiles/prestissimo-runtime.dockerfile
+++ b/presto-native-execution/scripts/dockerfiles/prestissimo-runtime.dockerfile
@@ -26,8 +26,10 @@ ENV BUILD_DIR=""
 
 RUN mkdir -p /prestissimo /runtime-libraries
 COPY . /prestissimo/
-RUN EXTRA_CMAKE_FLAGS=${EXTRA_CMAKE_FLAGS} \
-    NUM_THREADS=${NUM_THREADS} make --directory="/prestissimo/" cmake-and-build BUILD_TYPE=${BUILD_TYPE} BUILD_DIR=${BUILD_DIR} BUILD_BASE_DIR=${BUILD_BASE_DIR}
+RUN --mount=type=cache,target=/root/.ccache,sharing=locked \
+    EXTRA_CMAKE_FLAGS=${EXTRA_CMAKE_FLAGS} \
+    NUM_THREADS=${NUM_THREADS} make --directory="/prestissimo/" cmake-and-build BUILD_TYPE=${BUILD_TYPE} BUILD_DIR=${BUILD_DIR} BUILD_BASE_DIR=${BUILD_BASE_DIR} && \
+    ccache -sz -v
 RUN !(LD_LIBRARY_PATH=${LD_LIBRARY_PATH}:/usr/local/lib:/usr/local/lib64 ldd /prestissimo/${BUILD_BASE_DIR}/${BUILD_DIR}/presto_cpp/main/presto_server  | grep "not found") && \
     LD_LIBRARY_PATH=${LD_LIBRARY_PATH}:/usr/local/lib:/usr/local/lib64 ldd /prestissimo/${BUILD_BASE_DIR}/${BUILD_DIR}/presto_cpp/main/presto_server | awk 'NF == 4 { system("cp " $3 " /runtime-libraries") }'
 
diff --git a/presto-native-execution/scripts/dockerfiles/ubuntu-22.04-dependency.dockerfile b/presto-native-execution/scripts/dockerfiles/ubuntu-22.04-dependency.dockerfile
index 2738552101ba3..8c548c2679adb 100644
--- a/presto-native-execution/scripts/dockerfiles/ubuntu-22.04-dependency.dockerfile
+++ b/presto-native-execution/scripts/dockerfiles/ubuntu-22.04-dependency.dockerfile
@@ -16,12 +16,18 @@ FROM ${base}
 # Set a default timezone, can be overriden via ARG
 ARG tz="America/New_York"
 ARG DEBIAN_FRONTEND="noninteractive"
-ENV PROMPT_ALWAYS_RESPOND=n
+
+# Set this when build arm with common flags
+# from https://github.com/facebookincubator/velox/pull/14366
+ARG ARM_BUILD_TARGET
+
+ENV PROMPT_ALWAYS_RESPOND=y
 ENV SUDO=" "
 # TZ and DEBIAN_FRONTEND="noninteractive"
 # are required to avoid tzdata installation
 # to prompt for region selection.
 ENV TZ=${tz}
+ENV ARM_BUILD_TARGET=${ARM_BUILD_TARGET}
 
 RUN mkdir -p /scripts /velox/scripts
 COPY scripts /scripts
diff --git a/presto-native-execution/scripts/setup-adapters.sh b/presto-native-execution/scripts/setup-adapters.sh
index 3cb965fe71781..42596d1c059b1 100755
--- a/presto-native-execution/scripts/setup-adapters.sh
+++ b/presto-native-execution/scripts/setup-adapters.sh
@@ -15,14 +15,21 @@
 set -eufx -o pipefail
 
 SCRIPT_DIR=$(readlink -f "$(dirname "${BASH_SOURCE[0]}")")
-if [ -f "${SCRIPT_DIR}/setup-helper-functions.sh" ]
+if [ -f "${SCRIPT_DIR}/setup-common.sh" ]
 then
-  source "${SCRIPT_DIR}/setup-helper-functions.sh"
+  source "${SCRIPT_DIR}/setup-common.sh"
 else
-  source "${SCRIPT_DIR}/../velox/scripts/setup-helper-functions.sh"
+  source "${SCRIPT_DIR}/../velox/scripts/setup-common.sh"
 fi
 DEPENDENCY_DIR=${DEPENDENCY_DIR:-$(pwd)}
 
+OS=$(uname)
+if [ "$OS" = "Darwin" ]; then
+  export INSTALL_PREFIX=${INSTALL_PREFIX:-"$(pwd)/deps-install"}
+else
+  export INSTALL_PREFIX=${INSTALL_PREFIX:-"/usr/local"}
+fi
+
 function install_jwt_cpp {
   github_checkout Thalhammer/jwt-cpp v0.6.0 --depth 1
   cmake_install -DBUILD_TESTS=OFF -DJWT_BUILD_EXAMPLES=OFF -DJWT_DISABLE_PICOJSON=ON -DJWT_CMAKE_FILES_INSTALL_DIR="${DEPENDENCY_DIR}/jwt-cpp"
@@ -35,62 +42,13 @@ function install_prometheus_cpp {
   cmake_install -DBUILD_SHARED_LIBS=ON -DENABLE_PUSH=OFF -DENABLE_COMPRESSION=OFF
 }
 
-function install_abseil {
-  # abseil-cpp
-  github_checkout abseil/abseil-cpp 20240116.2 --depth 1
-  cmake_install \
-    -DABSL_BUILD_TESTING=OFF \
-    -DCMAKE_CXX_STANDARD=17 \
-    -DABSL_PROPAGATE_CXX_STD=ON \
-    -DABSL_ENABLE_INSTALL=ON
-}
-
-function install_grpc {
-  # grpc
-  github_checkout grpc/grpc v1.48.1 --depth 1
-  cmake_install \
-    -DgRPC_BUILD_TESTS=OFF \
-    -DgRPC_ABSL_PROVIDER=package \
-    -DgRPC_ZLIB_PROVIDER=package \
-    -DgRPC_CARES_PROVIDER=package \
-    -DgRPC_RE2_PROVIDER=package \
-    -DgRPC_SSL_PROVIDER=package \
-    -DgRPC_PROTOBUF_PROVIDER=package \
-    -DgRPC_INSTALL=ON
-}
-
 function install_arrow_flight {
-  ARROW_VERSION="${ARROW_VERSION:-15.0.0}"
-  if [[ "$OSTYPE" == "linux-gnu"* ]]; then
-    export INSTALL_PREFIX=${INSTALL_PREFIX:-"/usr/local"}
-    LINUX_DISTRIBUTION=$(. /etc/os-release && echo ${ID})
-    if [[ "$LINUX_DISTRIBUTION" == "ubuntu" || "$LINUX_DISTRIBUTION" == "debian" ]]; then
-      SUDO="${SUDO:-"sudo --preserve-env"}"
-      ${SUDO} apt install -y libc-ares-dev
-      ${SUDO} ldconfig -v 2>/dev/null | grep "${INSTALL_PREFIX}/lib" || \
-        echo "${INSTALL_PREFIX}/lib" | ${SUDO} tee /etc/ld.so.conf.d/local-libraries.conf > /dev/null \
-        && ${SUDO} ldconfig
-    else
-      dnf -y install c-ares-devel
-      ldconfig -v 2>/dev/null | grep "${INSTALL_PREFIX}/lib" || \
-        echo "${INSTALL_PREFIX}/lib" | tee /etc/ld.so.conf.d/local-libraries.conf > /dev/null \
-        && ldconfig
-    fi
-  else
-    # The installation script for the Arrow Flight connector currently works only on Linux distributions.
-    return 0
-  fi
-
-  install_abseil
-  install_grpc
-
+  # Velox provides an install for the Arrow library. Rebuild with the original Velox options and
+  # Arrow Flight enabled. The Velox version of Arrow is used.
   # NOTE: benchmarks are on due to a compilation error with v15.0.0, once updated that can be removed
   # see https://github.com/apache/arrow/issues/41617
-  wget_and_untar https://github.com/apache/arrow/archive/apache-arrow-${ARROW_VERSION}.tar.gz arrow
-  cmake_install_dir arrow/cpp \
-    -DARROW_FLIGHT=ON \
-    -DARROW_BUILD_BENCHMARKS=ON \
-    -DCMAKE_INSTALL_PREFIX=${INSTALL_PREFIX}
+  EXTRA_ARROW_OPTIONS=" -DARROW_FLIGHT=ON -DARROW_BUILD_BENCHMARKS=ON -DgRPC_SOURCE=BUNDLED -DProtobuf_SOURCE=BUNDLED "
+  install_arrow
 }
 
 cd "${DEPENDENCY_DIR}" || exit
diff --git a/presto-native-execution/src/test/java/com/facebook/presto/nativeworker/AbstractTestExpressionCompiler.java b/presto-native-execution/src/test/java/com/facebook/presto/nativeworker/AbstractTestExpressionCompiler.java
index 9ed6cbf0e9981..dec795a4b6fca 100644
--- a/presto-native-execution/src/test/java/com/facebook/presto/nativeworker/AbstractTestExpressionCompiler.java
+++ b/presto-native-execution/src/test/java/com/facebook/presto/nativeworker/AbstractTestExpressionCompiler.java
@@ -84,12 +84,80 @@ public void testNestedColumnFilter()
     {
     }
 
+    // Caused by: com.facebook.presto.spi.PrestoException: Catalog does not exist: hive
     @Override
     @Ignore
+    public void testBinaryOperatorsBigintDecimal()
+    {
+    }
+
+    // Caused by: com.facebook.presto.spi.PrestoException: Catalog does not exist: hive
+    @Override
+    @Ignore
+    public void testBinaryOperatorsBoolean()
+    {
+    }
+
+    // Caused by: com.facebook.presto.spi.PrestoException: Catalog does not exist: hive
+    @Override
+    @Ignore
+    public void testBinaryOperatorsDoubleDecimal()
+    {
+    }
+
+    // Caused by: com.facebook.presto.spi.PrestoException: Catalog does not exist: hive
+    @Override
+    @Ignore
+    public void testBinaryOperatorsDoubleDouble()
+    {
+    }
+
     public void testTernaryOperatorsLongLong()
     {
     }
 
+    // Caused by: com.facebook.presto.spi.PrestoException: Catalog does not exist: hive
+    @Override
+    @Ignore
+    public void testBinaryOperatorsDoubleIntegral()
+    {
+    }
+
+    // Caused by: com.facebook.presto.spi.PrestoException: Catalog does not exist: hive
+    @Override
+    @Ignore
+    public void testBinaryOperatorsIntegerDecimal()
+    {
+    }
+
+    // Caused by: com.facebook.presto.spi.PrestoException: Catalog does not exist: hive
+    @Override
+    @Ignore
+    public void testBinaryOperatorsIntegralDouble()
+    {
+    }
+
+    // Caused by: com.facebook.presto.spi.PrestoException: Catalog does not exist: hive
+    @Override
+    @Ignore
+    public void testBinaryOperatorsIntegralIntegral()
+    {
+    }
+
+    // Caused by: com.facebook.presto.spi.PrestoException: Catalog does not exist: hive
+    @Override
+    @Ignore
+    public void testBinaryOperatorsString()
+    {
+    }
+
+    // Caused by: com.facebook.presto.spi.PrestoException: Catalog does not exist: hive
+    @Override
+    @Ignore
+    public void testNullif()
+    {
+    }
+
     @Override
     @Ignore
     public void testTernaryOperatorsLongDouble()
diff --git a/presto-native-execution/src/test/java/com/facebook/presto/nativeworker/AbstractTestNativeGeneralQueries.java b/presto-native-execution/src/test/java/com/facebook/presto/nativeworker/AbstractTestNativeGeneralQueries.java
index e57585fd06eee..1549ccddbaaef 100644
--- a/presto-native-execution/src/test/java/com/facebook/presto/nativeworker/AbstractTestNativeGeneralQueries.java
+++ b/presto-native-execution/src/test/java/com/facebook/presto/nativeworker/AbstractTestNativeGeneralQueries.java
@@ -1211,7 +1211,7 @@ public void testSetSessionJavaWorkerSessionProperty()
                 "MaterializedResult{rows=[[true]], " +
                         "types=[boolean], " +
                         "setSessionProperties={distinct_aggregation_spill_enabled=false}, " +
-                        "resetSessionProperties=[], updateType=SET SESSION}");
+                        "resetSessionProperties=[], updateType=SET SESSION, clearTransactionId=false}");
     }
 
     @Test
diff --git a/presto-native-execution/src/test/java/com/facebook/presto/nativeworker/NativeQueryRunnerUtils.java b/presto-native-execution/src/test/java/com/facebook/presto/nativeworker/NativeQueryRunnerUtils.java
index 3650f8cb36531..59a60b0024842 100644
--- a/presto-native-execution/src/test/java/com/facebook/presto/nativeworker/NativeQueryRunnerUtils.java
+++ b/presto-native-execution/src/test/java/com/facebook/presto/nativeworker/NativeQueryRunnerUtils.java
@@ -29,7 +29,7 @@ private NativeQueryRunnerUtils() {}
     public static Map getNativeWorkerHiveProperties()
     {
         return ImmutableMap.of("hive.parquet.pushdown-filter-enabled", "true",
-            "hive.orc-compression-codec", "ZSTD", "hive.storage-format", "DWRF");
+                "hive.orc-compression-codec", "ZSTD", "hive.storage-format", "DWRF");
     }
 
     public static Map getNativeWorkerIcebergProperties()
@@ -59,6 +59,8 @@ public static Map getNativeSidecarProperties()
                 .put("coordinator-sidecar-enabled", "true")
                 .put("exclude-invalid-worker-session-properties", "true")
                 .put("presto.default-namespace", "native.default")
+                // inline-sql-functions is overridden to be true in sidecar enabled native clusters.
+                .put("inline-sql-functions", "true")
                 .build();
     }
 
diff --git a/presto-native-execution/src/test/java/com/facebook/presto/nativeworker/PrestoNativeQueryRunnerUtils.java b/presto-native-execution/src/test/java/com/facebook/presto/nativeworker/PrestoNativeQueryRunnerUtils.java
index 54980ba54782f..e287ef7e0cb20 100644
--- a/presto-native-execution/src/test/java/com/facebook/presto/nativeworker/PrestoNativeQueryRunnerUtils.java
+++ b/presto-native-execution/src/test/java/com/facebook/presto/nativeworker/PrestoNativeQueryRunnerUtils.java
@@ -123,9 +123,11 @@ public static class HiveQueryRunnerBuilder
         private String security;
         private boolean addStorageFormatToPath;
         private boolean coordinatorSidecarEnabled;
+        private boolean builtInWorkerFunctionsEnabled;
         private boolean enableRuntimeMetricsCollection;
         private boolean enableSsdCache;
         private boolean failOnNestedLoopJoin;
+        private boolean implicitCastCharNToVarchar;
         // External worker launcher is applicable only for the native hive query runner, since it depends on other
         // properties it should be created once all the other query runner configs are set. This variable indicates
         // whether the query runner returned by builder should use an external worker launcher, it will be true only
@@ -220,6 +222,12 @@ public HiveQueryRunnerBuilder setCoordinatorSidecarEnabled(boolean coordinatorSi
             return this;
         }
 
+        public HiveQueryRunnerBuilder setBuiltInWorkerFunctionsEnabled(boolean builtInWorkerFunctionsEnabled)
+        {
+            this.builtInWorkerFunctionsEnabled = builtInWorkerFunctionsEnabled;
+            return this;
+        }
+
         public HiveQueryRunnerBuilder setStorageFormat(String storageFormat)
         {
             this.storageFormat = storageFormat;
@@ -239,6 +247,12 @@ public HiveQueryRunnerBuilder setCacheMaxSize(Integer cacheMaxSize)
             return this;
         }
 
+        public HiveQueryRunnerBuilder setImplicitCastCharNToVarchar(boolean implicitCastCharNToVarchar)
+        {
+            this.implicitCastCharNToVarchar = implicitCastCharNToVarchar;
+            return this;
+        }
+
         public HiveQueryRunnerBuilder setExtraProperties(Map extraProperties)
         {
             this.extraProperties.putAll(extraProperties);
@@ -263,7 +277,7 @@ public QueryRunner build()
             Optional> externalWorkerLauncher = Optional.empty();
             if (this.useExternalWorkerLauncher) {
                 externalWorkerLauncher = getExternalWorkerLauncher("hive", serverBinary, cacheMaxSize, remoteFunctionServerUds,
-                        failOnNestedLoopJoin, coordinatorSidecarEnabled, enableRuntimeMetricsCollection, enableSsdCache);
+                        failOnNestedLoopJoin, coordinatorSidecarEnabled, builtInWorkerFunctionsEnabled, enableRuntimeMetricsCollection, enableSsdCache, implicitCastCharNToVarchar);
             }
             return HiveQueryRunner.createQueryRunner(
                     ImmutableList.of(),
@@ -352,7 +366,7 @@ public QueryRunner build()
             Optional> externalWorkerLauncher = Optional.empty();
             if (this.useExternalWorkerLauncher) {
                 externalWorkerLauncher = getExternalWorkerLauncher("iceberg", serverBinary, cacheMaxSize, remoteFunctionServerUds,
-                        false, false, false, false);
+                        false, false, false, false, false, false);
             }
             return IcebergQueryRunner.builder()
                     .setExtraProperties(extraProperties)
@@ -448,8 +462,10 @@ public static Optional> getExternalWorkerLaunc
             Optional remoteFunctionServerUds,
             Boolean failOnNestedLoopJoin,
             boolean isCoordinatorSidecarEnabled,
+            boolean isBuiltInWorkerFunctionsEnabled,
             boolean enableRuntimeMetricsCollection,
-            boolean enableSsdCache)
+            boolean enableSsdCache,
+            boolean implicitCastCharNToVarchar)
     {
         return
                 Optional.of((workerIndex, discoveryUri) -> {
@@ -470,6 +486,10 @@ public static Optional> getExternalWorkerLaunc
                                     "native-sidecar=true%n" +
                                     "presto.default-namespace=native.default%n", configProperties);
                         }
+                        else if (isBuiltInWorkerFunctionsEnabled) {
+                            configProperties = format("%s%n" +
+                                    "native-sidecar=true%n", configProperties);
+                        }
 
                         if (enableRuntimeMetricsCollection) {
                             configProperties = format("%s%n" +
@@ -497,6 +517,10 @@ public static Optional> getExternalWorkerLaunc
                             configProperties = format("%s%n" + "velox-plan-validator-fail-on-nested-loop-join=true%n", configProperties);
                         }
 
+                        if (implicitCastCharNToVarchar) {
+                            configProperties = format("%s%n" + "char-n-to-varchar-implicit-cast=true%n", configProperties);
+                        }
+
                         Files.write(tempDirectoryPath.resolve("config.properties"), configProperties.getBytes());
                         Files.write(tempDirectoryPath.resolve("node.properties"),
                                 format("node.id=%s%n" +
diff --git a/presto-native-execution/src/test/java/com/facebook/presto/nativeworker/TestBuiltInNativeFunctions.java b/presto-native-execution/src/test/java/com/facebook/presto/nativeworker/TestBuiltInNativeFunctions.java
new file mode 100644
index 0000000000000..f22707cf0bfad
--- /dev/null
+++ b/presto-native-execution/src/test/java/com/facebook/presto/nativeworker/TestBuiltInNativeFunctions.java
@@ -0,0 +1,178 @@
+/*
+ * Licensed 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 com.facebook.presto.nativeworker;
+
+import com.facebook.presto.cost.CostCalculator;
+import com.facebook.presto.cost.CostCalculatorUsingExchanges;
+import com.facebook.presto.cost.CostCalculatorWithEstimatedExchanges;
+import com.facebook.presto.cost.CostComparator;
+import com.facebook.presto.cost.TaskCountEstimator;
+import com.facebook.presto.execution.QueryManagerConfig;
+import com.facebook.presto.execution.TaskManagerConfig;
+import com.facebook.presto.metadata.InMemoryNodeManager;
+import com.facebook.presto.metadata.Metadata;
+import com.facebook.presto.nodeManager.PluginNodeManager;
+import com.facebook.presto.spi.WarningCollector;
+import com.facebook.presto.sql.analyzer.FeaturesConfig;
+import com.facebook.presto.sql.analyzer.QueryExplainer;
+import com.facebook.presto.sql.expressions.ExpressionOptimizerManager;
+import com.facebook.presto.sql.planner.PartitioningProviderManager;
+import com.facebook.presto.sql.planner.PlanFragmenter;
+import com.facebook.presto.sql.planner.PlanOptimizers;
+import com.facebook.presto.sql.planner.optimizations.PlanOptimizer;
+import com.facebook.presto.sql.planner.sanity.PlanChecker;
+import com.facebook.presto.sql.tree.ExplainType;
+import com.facebook.presto.testing.QueryRunner;
+import com.facebook.presto.tests.AbstractTestQueryFramework;
+import com.facebook.presto.tests.DistributedQueryRunner;
+import com.google.common.collect.ImmutableMap;
+import org.intellij.lang.annotations.Language;
+import org.testng.annotations.Test;
+import org.weakref.jmx.MBeanExporter;
+import org.weakref.jmx.testing.TestingMBeanServer;
+
+import java.util.List;
+import java.util.regex.Pattern;
+
+import static com.facebook.presto.nativeworker.NativeQueryRunnerUtils.createLineitem;
+import static com.facebook.presto.nativeworker.NativeQueryRunnerUtils.createNation;
+import static com.facebook.presto.nativeworker.NativeQueryRunnerUtils.createOrders;
+import static com.facebook.presto.nativeworker.NativeQueryRunnerUtils.createOrdersEx;
+import static com.facebook.presto.nativeworker.NativeQueryRunnerUtils.createRegion;
+import static com.facebook.presto.transaction.TransactionBuilder.transaction;
+import static com.facebook.presto.util.AnalyzerUtil.createParsingOptions;
+import static java.util.Collections.emptyList;
+import static org.testng.Assert.fail;
+
+public class TestBuiltInNativeFunctions
+        extends AbstractTestQueryFramework
+{
+    @Override
+    protected void createTables()
+    {
+        QueryRunner queryRunner = (QueryRunner) getExpectedQueryRunner();
+        createLineitem(queryRunner);
+        createNation(queryRunner);
+        createOrders(queryRunner);
+        createOrdersEx(queryRunner);
+        createRegion(queryRunner);
+    }
+
+    @Override
+    protected QueryRunner createQueryRunner()
+            throws Exception
+    {
+        DistributedQueryRunner queryRunner = (DistributedQueryRunner) PrestoNativeQueryRunnerUtils.nativeHiveQueryRunnerBuilder()
+                .setExtraProperties(ImmutableMap.of("built-in-sidecar-functions-enabled", "true"))
+                .setAddStorageFormatToPath(true)
+                .setBuiltInWorkerFunctionsEnabled(true)
+                .build();
+
+        queryRunner.registerNativeFunctions();
+
+        return queryRunner;
+    }
+
+    @Override
+    protected QueryRunner createExpectedQueryRunner()
+            throws Exception
+    {
+        return PrestoNativeQueryRunnerUtils.javaHiveQueryRunnerBuilder()
+                .setAddStorageFormatToPath(true)
+                .build();
+    }
+
+    private void assertJsonPlan(@Language("SQL") String query, boolean withBuiltInSidecarEnabled, @Language("RegExp") String jsonPlanRegex, boolean shouldContainRegex)
+    {
+        QueryRunner queryRunner;
+        if (withBuiltInSidecarEnabled) {
+            queryRunner = getQueryRunner();
+        }
+        else {
+            queryRunner = (QueryRunner) getExpectedQueryRunner();
+        }
+
+        QueryExplainer explainer = getQueryExplainerFromProvidedQueryRunner(queryRunner);
+        transaction(queryRunner.getTransactionManager(), queryRunner.getAccessControl())
+                .singleStatement()
+                .execute(queryRunner.getDefaultSession(), transactionSession -> {
+                    String actualPlan = explainer.getJsonPlan(transactionSession, getSqlParser().createStatement(query, createParsingOptions(transactionSession)), ExplainType.Type.LOGICAL, emptyList(), WarningCollector.NOOP, query);
+                    Pattern p = Pattern.compile(jsonPlanRegex, Pattern.MULTILINE);
+                    if (shouldContainRegex) {
+                        if (!p.matcher(actualPlan).find()) {
+                            fail("Query plan text does not contain regex");
+                        }
+                    }
+                    else {
+                        if (p.matcher(actualPlan).find()) {
+                            fail("Query plan text contains bad pattern");
+                        }
+                    }
+
+                    return null;
+                });
+    }
+
+    private QueryExplainer getQueryExplainerFromProvidedQueryRunner(QueryRunner queryRunner)
+    {
+        Metadata metadata = queryRunner.getMetadata();
+        FeaturesConfig featuresConfig = createFeaturesConfig();
+        boolean noExchange = queryRunner.getNodeCount() == 1;
+        TaskCountEstimator taskCountEstimator = new TaskCountEstimator(queryRunner::getNodeCount);
+        CostCalculator costCalculator = new CostCalculatorUsingExchanges(taskCountEstimator);
+        List optimizers = new PlanOptimizers(
+                metadata,
+                getSqlParser(),
+                noExchange,
+                new MBeanExporter(new TestingMBeanServer()),
+                queryRunner.getSplitManager(),
+                queryRunner.getPlanOptimizerManager(),
+                queryRunner.getPageSourceManager(),
+                queryRunner.getStatsCalculator(),
+                costCalculator,
+                new CostCalculatorWithEstimatedExchanges(costCalculator, taskCountEstimator),
+                new CostComparator(featuresConfig),
+                taskCountEstimator,
+                new PartitioningProviderManager(),
+                featuresConfig,
+                new ExpressionOptimizerManager(
+                        new PluginNodeManager(new InMemoryNodeManager()),
+                        queryRunner.getMetadata().getFunctionAndTypeManager()),
+                new TaskManagerConfig())
+                .getPlanningTimeOptimizers();
+        return new QueryExplainer(
+                optimizers,
+                new PlanFragmenter(metadata, queryRunner.getNodePartitioningManager(), new QueryManagerConfig(), featuresConfig, queryRunner.getPlanCheckerProviderManager()),
+                metadata,
+                queryRunner.getAccessControl(),
+                getSqlParser(),
+                queryRunner.getStatsCalculator(),
+                costCalculator,
+                ImmutableMap.of(),
+                new PlanChecker(featuresConfig, false, queryRunner.getPlanCheckerProviderManager()));
+    }
+
+    @Test
+    public void testUdfQueries()
+    {
+        assertQuery("SELECT ARRAY['abc']");
+        assertQuery("SELECT ARRAY[1, 2, 3]");
+        assertQuery("SELECT map_remove_null_values( MAP( ARRAY['a', 'b', 'c'], ARRAY[1, NULL, 3] ) )");
+        assertQuery("SELECT presto.default.map_remove_null_values( MAP( ARRAY['a', 'b', 'c'], ARRAY[1, NULL, 3] ) )");
+        assertQueryFails("SELECT native.default.map_remove_null_values( MAP( ARRAY['a', 'b', 'c'], ARRAY[1, NULL, 3] ) )", ".*Function native.default.map_remove_null_values not registered.*");
+        assertJsonPlan("SELECT map_remove_null_values( MAP( ARRAY['a', 'b', 'c'], ARRAY[1, NULL, 3] ) )", true, "lambda", false);
+        assertJsonPlan("SELECT map_remove_null_values( MAP( ARRAY['a', 'b', 'c'], ARRAY[1, NULL, 3] ) )", false, "lambda", true);
+    }
+}
diff --git a/presto-native-execution/src/test/java/com/facebook/presto/nativeworker/TestPrestoNativeArrayFunctionQueries.java b/presto-native-execution/src/test/java/com/facebook/presto/nativeworker/TestPrestoNativeArrayFunctionQueries.java
index bfa21641206b2..f45f20c545d00 100644
--- a/presto-native-execution/src/test/java/com/facebook/presto/nativeworker/TestPrestoNativeArrayFunctionQueries.java
+++ b/presto-native-execution/src/test/java/com/facebook/presto/nativeworker/TestPrestoNativeArrayFunctionQueries.java
@@ -13,6 +13,7 @@
  */
 package com.facebook.presto.nativeworker;
 
+import com.facebook.presto.scalar.sql.SqlInvokedFunctionsPlugin;
 import com.facebook.presto.testing.ExpectedQueryRunner;
 import com.facebook.presto.testing.QueryRunner;
 
@@ -20,18 +21,24 @@ public class TestPrestoNativeArrayFunctionQueries
         extends AbstractTestNativeArrayFunctionQueries
 {
     @Override
-    protected QueryRunner createQueryRunner() throws Exception
+    protected QueryRunner createQueryRunner()
+            throws Exception
     {
-        return PrestoNativeQueryRunnerUtils.nativeHiveQueryRunnerBuilder()
+        QueryRunner queryRunner = PrestoNativeQueryRunnerUtils.nativeHiveQueryRunnerBuilder()
                 .setAddStorageFormatToPath(true)
                 .build();
+        queryRunner.installPlugin(new SqlInvokedFunctionsPlugin());
+        return queryRunner;
     }
 
     @Override
-    protected ExpectedQueryRunner createExpectedQueryRunner() throws Exception
+    protected ExpectedQueryRunner createExpectedQueryRunner()
+            throws Exception
     {
-        return PrestoNativeQueryRunnerUtils.javaHiveQueryRunnerBuilder()
+        QueryRunner queryRunner = PrestoNativeQueryRunnerUtils.javaHiveQueryRunnerBuilder()
                 .setAddStorageFormatToPath(true)
                 .build();
+        queryRunner.installPlugin(new SqlInvokedFunctionsPlugin());
+        return queryRunner;
     }
 }
diff --git a/presto-native-execution/src/test/java/com/facebook/presto/nativeworker/TestPrestoNativeGeneralQueriesJSON.java b/presto-native-execution/src/test/java/com/facebook/presto/nativeworker/TestPrestoNativeGeneralQueriesJSON.java
index 4b1d27f30426e..bcba83d00da65 100644
--- a/presto-native-execution/src/test/java/com/facebook/presto/nativeworker/TestPrestoNativeGeneralQueriesJSON.java
+++ b/presto-native-execution/src/test/java/com/facebook/presto/nativeworker/TestPrestoNativeGeneralQueriesJSON.java
@@ -13,6 +13,7 @@
  */
 package com.facebook.presto.nativeworker;
 
+import com.facebook.presto.scalar.sql.SqlInvokedFunctionsPlugin;
 import com.facebook.presto.testing.ExpectedQueryRunner;
 import com.facebook.presto.testing.QueryRunner;
 
@@ -23,17 +24,21 @@ public class TestPrestoNativeGeneralQueriesJSON
     protected QueryRunner createQueryRunner()
             throws Exception
     {
-        return PrestoNativeQueryRunnerUtils.nativeHiveQueryRunnerBuilder()
+        QueryRunner queryRunner = PrestoNativeQueryRunnerUtils.nativeHiveQueryRunnerBuilder()
                 .setAddStorageFormatToPath(true)
                 .build();
+        queryRunner.installPlugin(new SqlInvokedFunctionsPlugin());
+        return queryRunner;
     }
 
     @Override
     protected ExpectedQueryRunner createExpectedQueryRunner()
             throws Exception
     {
-        return PrestoNativeQueryRunnerUtils.javaHiveQueryRunnerBuilder()
+        QueryRunner queryRunner = PrestoNativeQueryRunnerUtils.javaHiveQueryRunnerBuilder()
                 .setAddStorageFormatToPath(true)
                 .build();
+        queryRunner.installPlugin(new SqlInvokedFunctionsPlugin());
+        return queryRunner;
     }
 }
diff --git a/presto-native-execution/src/test/java/com/facebook/presto/nativeworker/TestPrestoNativeGeneralQueriesThrift.java b/presto-native-execution/src/test/java/com/facebook/presto/nativeworker/TestPrestoNativeGeneralQueriesThrift.java
index a1c33cd09f51c..fd82dcde9209e 100644
--- a/presto-native-execution/src/test/java/com/facebook/presto/nativeworker/TestPrestoNativeGeneralQueriesThrift.java
+++ b/presto-native-execution/src/test/java/com/facebook/presto/nativeworker/TestPrestoNativeGeneralQueriesThrift.java
@@ -13,6 +13,7 @@
  */
 package com.facebook.presto.nativeworker;
 
+import com.facebook.presto.scalar.sql.SqlInvokedFunctionsPlugin;
 import com.facebook.presto.testing.ExpectedQueryRunner;
 import com.facebook.presto.testing.QueryRunner;
 
@@ -23,18 +24,22 @@ public class TestPrestoNativeGeneralQueriesThrift
     protected QueryRunner createQueryRunner()
             throws Exception
     {
-        return PrestoNativeQueryRunnerUtils.nativeHiveQueryRunnerBuilder()
+        QueryRunner queryRunner = PrestoNativeQueryRunnerUtils.nativeHiveQueryRunnerBuilder()
                 .setAddStorageFormatToPath(true)
                 .setUseThrift(true)
                 .build();
+        queryRunner.installPlugin(new SqlInvokedFunctionsPlugin());
+        return queryRunner;
     }
 
     @Override
     protected ExpectedQueryRunner createExpectedQueryRunner()
             throws Exception
     {
-        return PrestoNativeQueryRunnerUtils.javaHiveQueryRunnerBuilder()
+        QueryRunner queryRunner = PrestoNativeQueryRunnerUtils.javaHiveQueryRunnerBuilder()
                 .setAddStorageFormatToPath(true)
                 .build();
+        queryRunner.installPlugin(new SqlInvokedFunctionsPlugin());
+        return queryRunner;
     }
 }
diff --git a/presto-native-execution/src/test/java/com/facebook/presto/spark/PrestoSparkNativeQueryRunnerUtils.java b/presto-native-execution/src/test/java/com/facebook/presto/spark/PrestoSparkNativeQueryRunnerUtils.java
index f51acf99a6155..a0856f8c6a530 100644
--- a/presto-native-execution/src/test/java/com/facebook/presto/spark/PrestoSparkNativeQueryRunnerUtils.java
+++ b/presto-native-execution/src/test/java/com/facebook/presto/spark/PrestoSparkNativeQueryRunnerUtils.java
@@ -18,7 +18,6 @@
 import com.facebook.presto.hive.metastore.ExtendedHiveMetastore;
 import com.facebook.presto.nativeworker.PrestoNativeQueryRunnerUtils;
 import com.facebook.presto.spark.execution.nativeprocess.NativeExecutionModule;
-import com.facebook.presto.spark.execution.property.NativeExecutionConnectorConfig;
 import com.facebook.presto.spi.security.PrincipalType;
 import com.facebook.presto.testing.QueryRunner;
 import com.google.common.collect.ImmutableList;
@@ -80,7 +79,6 @@ public static Map getNativeExecutionSparkConfigs()
                 .put("catalog.config-dir", "/")
                 .put("task.info-update-interval", "100ms")
                 .put("spark.initial-partition-count", "1")
-                .put("register-test-functions", "true")
                 .put("native-execution-program-arguments", "--logtostderr=1 --minloglevel=3")
                 .put("spark.partition-count-auto-tune-enabled", "false");
 
@@ -101,7 +99,8 @@ public static Map getNativeExecutionSparkConfigs()
 
     public static PrestoSparkQueryRunner createHiveRunner()
     {
-        PrestoSparkQueryRunner queryRunner = createRunner("hive", new NativeExecutionModule());
+        PrestoSparkQueryRunner queryRunner = createRunner("hive", new NativeExecutionModule(
+                Optional.of(ImmutableMap.of("hive", ImmutableMap.of("connector.name", "hive")))));
         PrestoNativeQueryRunnerUtils.setupJsonFunctionNamespaceManager(queryRunner, "external_functions.json", "json");
 
         return queryRunner;
@@ -125,7 +124,8 @@ public static PrestoSparkQueryRunner createTpchRunner()
         return createRunner(
                 "tpchstandard",
                 new NativeExecutionModule(
-                        Optional.of(new NativeExecutionConnectorConfig().setConnectorName("tpch"))));
+                        Optional.of(
+                                ImmutableMap.of("hive", ImmutableMap.of("connector.name", "tpch")))));
     }
 
     public static PrestoSparkQueryRunner createRunner(String defaultCatalog, Optional baseDir, Map additionalConfigProperties, Map additionalSparkProperties, ImmutableList nativeModules)
@@ -187,6 +187,7 @@ private static Database createDatabaseMetastoreObject(String name)
     private static Map getNativeExecutionShuffleConfigs()
     {
         ImmutableMap.Builder sparkConfigs = ImmutableMap.builder();
+        sparkConfigs.put("spark.ui.enabled", "false");
         sparkConfigs.put(SPARK_SHUFFLE_MANAGER, "com.facebook.presto.spark.classloader_interface.PrestoSparkNativeExecutionShuffleManager");
         sparkConfigs.put(FALLBACK_SPARK_SHUFFLE_MANAGER, "org.apache.spark.shuffle.sort.SortShuffleManager");
         return sparkConfigs.build();
diff --git a/presto-native-execution/src/test/java/com/facebook/presto/spark/TestPrestoSparkNativeArrayFunctionQueries.java b/presto-native-execution/src/test/java/com/facebook/presto/spark/TestPrestoSparkNativeArrayFunctionQueries.java
index 40db300489142..f513715b25f79 100644
--- a/presto-native-execution/src/test/java/com/facebook/presto/spark/TestPrestoSparkNativeArrayFunctionQueries.java
+++ b/presto-native-execution/src/test/java/com/facebook/presto/spark/TestPrestoSparkNativeArrayFunctionQueries.java
@@ -14,6 +14,7 @@
 package com.facebook.presto.spark;
 
 import com.facebook.presto.nativeworker.AbstractTestNativeArrayFunctionQueries;
+import com.facebook.presto.scalar.sql.SqlInvokedFunctionsPlugin;
 import com.facebook.presto.testing.ExpectedQueryRunner;
 import com.facebook.presto.testing.QueryRunner;
 
@@ -23,13 +24,21 @@ public class TestPrestoSparkNativeArrayFunctionQueries
     @Override
     protected QueryRunner createQueryRunner()
     {
-        return PrestoSparkNativeQueryRunnerUtils.createHiveRunner();
+        QueryRunner queryRunner = PrestoSparkNativeQueryRunnerUtils.createHiveRunner();
+
+        // Install plugins needed for extra array functions.
+        queryRunner.installPlugin(new SqlInvokedFunctionsPlugin());
+        return queryRunner;
     }
 
     @Override
     protected ExpectedQueryRunner createExpectedQueryRunner()
             throws Exception
     {
-        return PrestoSparkNativeQueryRunnerUtils.createJavaQueryRunner();
+        QueryRunner queryRunner = PrestoSparkNativeQueryRunnerUtils.createJavaQueryRunner();
+
+        // Install plugins needed for extra array functions.
+        queryRunner.installPlugin(new SqlInvokedFunctionsPlugin());
+        return queryRunner;
     }
 }
diff --git a/presto-native-execution/src/test/java/com/facebook/presto/spark/TestPrestoSparkNativeGeneralQueries.java b/presto-native-execution/src/test/java/com/facebook/presto/spark/TestPrestoSparkNativeGeneralQueries.java
index 902de47b2fd79..d67dba20dfa93 100644
--- a/presto-native-execution/src/test/java/com/facebook/presto/spark/TestPrestoSparkNativeGeneralQueries.java
+++ b/presto-native-execution/src/test/java/com/facebook/presto/spark/TestPrestoSparkNativeGeneralQueries.java
@@ -16,6 +16,7 @@
 import com.facebook.airlift.log.Level;
 import com.facebook.airlift.log.Logging;
 import com.facebook.presto.nativeworker.AbstractTestNativeGeneralQueries;
+import com.facebook.presto.scalar.sql.SqlInvokedFunctionsPlugin;
 import com.facebook.presto.testing.ExpectedQueryRunner;
 import com.facebook.presto.testing.QueryRunner;
 import org.testng.annotations.Ignore;
@@ -31,14 +32,22 @@ public class TestPrestoSparkNativeGeneralQueries
     @Override
     protected QueryRunner createQueryRunner()
     {
-        return PrestoSparkNativeQueryRunnerUtils.createHiveRunner();
+        QueryRunner queryRunner = PrestoSparkNativeQueryRunnerUtils.createHiveRunner();
+
+        // Install plugins needed for extra array functions.
+        queryRunner.installPlugin(new SqlInvokedFunctionsPlugin());
+        return queryRunner;
     }
 
     @Override
     protected ExpectedQueryRunner createExpectedQueryRunner()
             throws Exception
     {
-        return PrestoSparkNativeQueryRunnerUtils.createJavaQueryRunner();
+        QueryRunner queryRunner = PrestoSparkNativeQueryRunnerUtils.createJavaQueryRunner();
+
+        // Install plugins needed for extra array functions.
+        queryRunner.installPlugin(new SqlInvokedFunctionsPlugin());
+        return queryRunner;
     }
 
     @Override
@@ -106,4 +115,9 @@ public void testRowWiseExchange() {}
     @Override
     @Ignore
     public void testAnalyzeStatsOnDecimals() {}
+
+    // VeloxRuntimeError: it != connectors().end() Connector with ID 'hivecached' not registered
+    @Override
+    @Ignore
+    public void testCatalogWithCacheEnabled() {}
 }
diff --git a/presto-native-execution/src/test/java/com/facebook/presto/spark/TestPrestoSparkNativeTpchConnectorQueries.java b/presto-native-execution/src/test/java/com/facebook/presto/spark/TestPrestoSparkNativeTpchConnectorQueries.java
index b2a0bb1ad5d51..05d194f6f6754 100644
--- a/presto-native-execution/src/test/java/com/facebook/presto/spark/TestPrestoSparkNativeTpchConnectorQueries.java
+++ b/presto-native-execution/src/test/java/com/facebook/presto/spark/TestPrestoSparkNativeTpchConnectorQueries.java
@@ -43,4 +43,9 @@ public void testMissingTpchConnector()
     @Override
     @Ignore
     public void testTpchTinyTables() {}
+
+    // VeloxRuntimeError: it != connectors().end() Connector with ID 'tpchstandard' not registered
+    @Override
+    @Ignore
+    public void testTpchDateFilter() {}
 }
diff --git a/presto-native-execution/velox b/presto-native-execution/velox
index 8d01456cf77a5..71a310307a734 160000
--- a/presto-native-execution/velox
+++ b/presto-native-execution/velox
@@ -1 +1 @@
-Subproject commit 8d01456cf77a56d56c371ecc9509c5ae111157d8
+Subproject commit 71a310307a734e170206163efb504afb7e6d834c
diff --git a/presto-native-sidecar-plugin/pom.xml b/presto-native-sidecar-plugin/pom.xml
index 153ad18295f8c..d04424440469a 100644
--- a/presto-native-sidecar-plugin/pom.xml
+++ b/presto-native-sidecar-plugin/pom.xml
@@ -260,6 +260,26 @@
                 
             
         
+
+        
+            com.facebook.presto
+            presto-built-in-worker-function-tools
+            ${project.version}
+        
+
+        
+            com.facebook.presto
+            presto-native-sql-invoked-functions-plugin
+            ${project.version}
+            test
+        
+
+        
+            com.facebook.presto
+            presto-sql-invoked-functions-plugin
+            ${project.version}
+            test
+        
     
 
     
diff --git a/presto-native-sidecar-plugin/src/main/java/com/facebook/presto/sidecar/functionNamespace/NativeFunctionDefinitionProvider.java b/presto-native-sidecar-plugin/src/main/java/com/facebook/presto/sidecar/functionNamespace/NativeFunctionDefinitionProvider.java
index a0cd355cfe175..677752dd4d2c7 100644
--- a/presto-native-sidecar-plugin/src/main/java/com/facebook/presto/sidecar/functionNamespace/NativeFunctionDefinitionProvider.java
+++ b/presto-native-sidecar-plugin/src/main/java/com/facebook/presto/sidecar/functionNamespace/NativeFunctionDefinitionProvider.java
@@ -14,26 +14,24 @@
 package com.facebook.presto.sidecar.functionNamespace;
 
 import com.facebook.airlift.http.client.HttpClient;
-import com.facebook.airlift.http.client.HttpUriBuilder;
 import com.facebook.airlift.http.client.Request;
 import com.facebook.airlift.json.JsonCodec;
 import com.facebook.airlift.log.Logger;
 import com.facebook.presto.functionNamespace.JsonBasedUdfFunctionMetadata;
 import com.facebook.presto.functionNamespace.UdfFunctionSignatureMap;
 import com.facebook.presto.sidecar.ForSidecarInfo;
-import com.facebook.presto.spi.Node;
 import com.facebook.presto.spi.NodeManager;
 import com.facebook.presto.spi.PrestoException;
 import com.google.common.annotations.VisibleForTesting;
 import com.google.common.collect.ImmutableMap;
 import com.google.inject.Inject;
 
-import java.net.URI;
 import java.util.List;
 import java.util.Map;
 
 import static com.facebook.airlift.http.client.JsonResponseHandler.createJsonResponseHandler;
 import static com.facebook.airlift.http.client.Request.Builder.prepareGet;
+import static com.facebook.presto.builtin.tools.NativeSidecarFunctionRegistryTool.getSidecarLocationOnStartup;
 import static com.facebook.presto.spi.StandardErrorCode.INVALID_ARGUMENTS;
 import static java.util.Objects.requireNonNull;
 
@@ -42,27 +40,29 @@ public class NativeFunctionDefinitionProvider
 {
     private static final Logger log = Logger.get(NativeFunctionDefinitionProvider.class);
     private final JsonCodec>> nativeFunctionSignatureMapJsonCodec;
-    private final NodeManager nodeManager;
     private final HttpClient httpClient;
-    private static final String FUNCTION_SIGNATURES_ENDPOINT = "/v1/functions";
+    private final NativeFunctionNamespaceManagerConfig config;
 
     @Inject
     public NativeFunctionDefinitionProvider(
             @ForSidecarInfo HttpClient httpClient,
             JsonCodec>> nativeFunctionSignatureMapJsonCodec,
-            NodeManager nodeManager)
+            NativeFunctionNamespaceManagerConfig config)
     {
         this.nativeFunctionSignatureMapJsonCodec =
                 requireNonNull(nativeFunctionSignatureMapJsonCodec, "nativeFunctionSignatureMapJsonCodec is null");
-        this.nodeManager = requireNonNull(nodeManager, "nodeManager is null");
-        this.httpClient = requireNonNull(httpClient, "typeManager is null");
+        this.httpClient = requireNonNull(httpClient, "httpClient is null");
+        this.config = requireNonNull(config, "config is null");
     }
 
     @Override
     public UdfFunctionSignatureMap getUdfDefinition(NodeManager nodeManager)
     {
         try {
-            Request request = prepareGet().setUri(getSidecarLocation()).build();
+            Request request =
+                    prepareGet().setUri(
+                            getSidecarLocationOnStartup(
+                                    nodeManager, config.getSidecarNumRetries(), config.getSidecarRetryDelay().toMillis())).build();
             Map> nativeFunctionSignatureMap = httpClient.execute(request, createJsonResponseHandler(nativeFunctionSignatureMapJsonCodec));
             return new UdfFunctionSignatureMap(ImmutableMap.copyOf(nativeFunctionSignatureMap));
         }
@@ -71,15 +71,6 @@ public UdfFunctionSignatureMap getUdfDefinition(NodeManager nodeManager)
         }
     }
 
-    private URI getSidecarLocation()
-    {
-        Node sidecarNode = nodeManager.getSidecarNode();
-        return HttpUriBuilder
-                .uriBuilderFrom(sidecarNode.getHttpUri())
-                .appendPath(FUNCTION_SIGNATURES_ENDPOINT)
-                .build();
-    }
-
     @VisibleForTesting
     public HttpClient getHttpClient()
     {
diff --git a/presto-native-sidecar-plugin/src/main/java/com/facebook/presto/sidecar/functionNamespace/NativeFunctionNamespaceManager.java b/presto-native-sidecar-plugin/src/main/java/com/facebook/presto/sidecar/functionNamespace/NativeFunctionNamespaceManager.java
index 11ef917f5522d..f1b7142a188ee 100644
--- a/presto-native-sidecar-plugin/src/main/java/com/facebook/presto/sidecar/functionNamespace/NativeFunctionNamespaceManager.java
+++ b/presto-native-sidecar-plugin/src/main/java/com/facebook/presto/sidecar/functionNamespace/NativeFunctionNamespaceManager.java
@@ -14,14 +14,10 @@
 package com.facebook.presto.sidecar.functionNamespace;
 
 import com.facebook.airlift.log.Logger;
-import com.facebook.presto.common.CatalogSchemaName;
 import com.facebook.presto.common.QualifiedObjectName;
-import com.facebook.presto.common.type.NamedTypeSignature;
-import com.facebook.presto.common.type.StandardTypes;
 import com.facebook.presto.common.type.Type;
 import com.facebook.presto.common.type.TypeManager;
 import com.facebook.presto.common.type.TypeSignature;
-import com.facebook.presto.common.type.TypeSignatureParameter;
 import com.facebook.presto.common.type.UserDefinedType;
 import com.facebook.presto.functionNamespace.AbstractSqlInvokedFunctionNamespaceManager;
 import com.facebook.presto.functionNamespace.JsonBasedUdfFunctionMetadata;
@@ -38,7 +34,6 @@
 import com.facebook.presto.spi.function.FunctionMetadata;
 import com.facebook.presto.spi.function.FunctionMetadataManager;
 import com.facebook.presto.spi.function.FunctionNamespaceTransactionHandle;
-import com.facebook.presto.spi.function.LongVariableConstraint;
 import com.facebook.presto.spi.function.Parameter;
 import com.facebook.presto.spi.function.ScalarFunctionImplementation;
 import com.facebook.presto.spi.function.Signature;
@@ -48,35 +43,29 @@
 import com.facebook.presto.spi.function.SqlFunctionSupplier;
 import com.facebook.presto.spi.function.SqlInvokedAggregationFunctionImplementation;
 import com.facebook.presto.spi.function.SqlInvokedFunction;
-import com.facebook.presto.spi.function.TypeVariableConstraint;
 import com.google.common.annotations.VisibleForTesting;
 import com.google.common.base.Suppliers;
 import com.google.common.cache.CacheBuilder;
 import com.google.common.cache.CacheLoader;
 import com.google.common.cache.LoadingCache;
-import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableMap;
 import com.google.common.util.concurrent.UncheckedExecutionException;
 import jakarta.inject.Inject;
 
-import java.util.ArrayList;
 import java.util.Collection;
-import java.util.Collections;
 import java.util.List;
 import java.util.Map;
 import java.util.Optional;
 import java.util.concurrent.ConcurrentHashMap;
 import java.util.function.Supplier;
 
+import static com.facebook.presto.builtin.tools.WorkerFunctionUtil.createSqlInvokedFunction;
 import static com.facebook.presto.common.type.TypeSignatureUtils.resolveIntermediateType;
 import static com.facebook.presto.spi.StandardErrorCode.DUPLICATE_FUNCTION_ERROR;
 import static com.facebook.presto.spi.StandardErrorCode.GENERIC_INTERNAL_ERROR;
 import static com.facebook.presto.spi.StandardErrorCode.GENERIC_USER_ERROR;
 import static com.facebook.presto.spi.StandardErrorCode.NOT_SUPPORTED;
-import static com.facebook.presto.spi.function.FunctionVersion.notVersioned;
-import static com.facebook.presto.spi.function.RoutineCharacteristics.Language.CPP;
 import static com.google.common.base.Preconditions.checkArgument;
-import static com.google.common.base.Preconditions.checkState;
 import static com.google.common.collect.ImmutableList.toImmutableList;
 import static com.google.common.collect.MoreCollectors.onlyElement;
 import static java.lang.String.format;
@@ -138,7 +127,7 @@ private synchronized void populateNamespaceManager(UdfFunctionSignatureMap udfFu
     {
         Map> udfSignatureMap = udfFunctionSignatureMap.getUDFSignatureMap();
         udfSignatureMap.forEach((name, metaInfoList) -> {
-            List functions = metaInfoList.stream().map(metaInfo -> createSqlInvokedFunction(name, metaInfo)).collect(toImmutableList());
+            List functions = metaInfoList.stream().map(metaInfo -> createSqlInvokedFunction(name, metaInfo, getCatalogName())).collect(toImmutableList());
             functions.forEach(this::createFunction);
         });
     }
@@ -200,44 +189,6 @@ private AggregationFunctionImplementation processSqlFunctionHandle(SqlFunctionHa
         return aggregationImplementationByHandle.get(sqlFunctionHandle);
     }
 
-    protected synchronized SqlInvokedFunction createSqlInvokedFunction(String functionName, JsonBasedUdfFunctionMetadata jsonBasedUdfFunctionMetaData)
-    {
-        checkState(jsonBasedUdfFunctionMetaData.getRoutineCharacteristics().getLanguage().equals(CPP), "NativeFunctionNamespaceManager only supports CPP UDF");
-        QualifiedObjectName qualifiedFunctionName = QualifiedObjectName.valueOf(new CatalogSchemaName(getCatalogName(), jsonBasedUdfFunctionMetaData.getSchema()), functionName);
-        List parameterNameList = jsonBasedUdfFunctionMetaData.getParamNames();
-        List parameterTypeList = convertApplicableTypeToVariable(jsonBasedUdfFunctionMetaData.getParamTypes());
-        List typeVariableConstraintsList = jsonBasedUdfFunctionMetaData.getTypeVariableConstraints().isPresent() ?
-                jsonBasedUdfFunctionMetaData.getTypeVariableConstraints().get() : Collections.emptyList();
-        List longVariableConstraintList = jsonBasedUdfFunctionMetaData.getLongVariableConstraints().isPresent() ?
-                jsonBasedUdfFunctionMetaData.getLongVariableConstraints().get() : Collections.emptyList();
-
-        TypeSignature outputType = convertApplicableTypeToVariable(jsonBasedUdfFunctionMetaData.getOutputType());
-        ImmutableList.Builder parameterBuilder = ImmutableList.builder();
-        for (int i = 0; i < parameterNameList.size(); i++) {
-            parameterBuilder.add(new Parameter(parameterNameList.get(i), parameterTypeList.get(i)));
-        }
-
-        Optional aggregationFunctionMetadata =
-                jsonBasedUdfFunctionMetaData.getAggregateMetadata()
-                        .map(metadata -> new AggregationFunctionMetadata(
-                                convertApplicableTypeToVariable(metadata.getIntermediateType()),
-                                metadata.isOrderSensitive()));
-
-        return new SqlInvokedFunction(
-                qualifiedFunctionName,
-                parameterBuilder.build(),
-                typeVariableConstraintsList,
-                longVariableConstraintList,
-                outputType,
-                jsonBasedUdfFunctionMetaData.getDocString(),
-                jsonBasedUdfFunctionMetaData.getRoutineCharacteristics(),
-                "",
-                jsonBasedUdfFunctionMetaData.getVariableArity(),
-                notVersioned(),
-                jsonBasedUdfFunctionMetaData.getFunctionKind(),
-                aggregationFunctionMetadata);
-    }
-
     @Override
     protected Collection fetchFunctionsDirect(QualifiedObjectName functionName)
     {
@@ -320,109 +271,12 @@ public final FunctionHandle getFunctionHandle(Optional typeSignaturesList = convertApplicableTypeToVariable(ImmutableList.of(typeSignature));
-        checkArgument(!typeSignaturesList.isEmpty(), "Type signature list is empty for : " + typeSignature);
-        return typeSignaturesList.get(0);
-    }
-
-    public static List convertApplicableTypeToVariable(List typeSignatures)
-    {
-        List newTypeSignaturesList = new ArrayList<>();
-        for (TypeSignature typeSignature : typeSignatures) {
-            if (!typeSignature.getParameters().isEmpty()) {
-                TypeSignature newTypeSignature =
-                        new TypeSignature(
-                                typeSignature.getBase(),
-                                getTypeSignatureParameters(
-                                        typeSignature,
-                                        typeSignature.getParameters()));
-                newTypeSignaturesList.add(newTypeSignature);
-            }
-            else {
-                newTypeSignaturesList.add(typeSignature);
-            }
-        }
-        return newTypeSignaturesList;
-    }
-
     @VisibleForTesting
     public FunctionDefinitionProvider getFunctionDefinitionProvider()
     {
         return functionDefinitionProvider;
     }
 
-    private static List getTypeSignatureParameters(
-            TypeSignature typeSignature,
-            List typeSignatureParameterList)
-    {
-        List newParameterTypeList = new ArrayList<>();
-        for (TypeSignatureParameter parameter : typeSignatureParameterList) {
-            if (parameter.isLongLiteral()) {
-                newParameterTypeList.add(parameter);
-                continue;
-            }
-
-            boolean isNamedTypeSignature = parameter.isNamedTypeSignature();
-            TypeSignature parameterTypeSignature;
-            // If it's a named type signatures only in the case of row signature types.
-            if (isNamedTypeSignature) {
-                parameterTypeSignature = parameter.getNamedTypeSignature().getTypeSignature();
-            }
-            else {
-                parameterTypeSignature = parameter.getTypeSignature();
-            }
-
-            if (parameterTypeSignature.getParameters().isEmpty()) {
-                boolean changeTypeToVariable = isDecimalTypeBase(typeSignature.getBase());
-                if (changeTypeToVariable) {
-                    newParameterTypeList.add(
-                            TypeSignatureParameter.of(parameterTypeSignature.getBase()));
-                }
-                else {
-                    if (isNamedTypeSignature) {
-                        newParameterTypeList.add(TypeSignatureParameter.of(parameter.getNamedTypeSignature()));
-                    }
-                    else {
-                        newParameterTypeList.add(TypeSignatureParameter.of(parameterTypeSignature));
-                    }
-                }
-            }
-            else {
-                TypeSignature newTypeSignature =
-                        new TypeSignature(
-                                parameterTypeSignature.getBase(),
-                                getTypeSignatureParameters(
-                                        parameterTypeSignature.getStandardTypeSignature(),
-                                        parameterTypeSignature.getParameters()));
-                if (isNamedTypeSignature) {
-                    newParameterTypeList.add(
-                            TypeSignatureParameter.of(
-                                    new NamedTypeSignature(
-                                            Optional.empty(),
-                                            newTypeSignature)));
-                }
-                else {
-                    newParameterTypeList.add(TypeSignatureParameter.of(newTypeSignature));
-                }
-            }
-        }
-        return newParameterTypeList;
-    }
-
-    private static boolean isDecimalTypeBase(String typeBase)
-    {
-        return typeBase.equals(StandardTypes.DECIMAL);
-    }
-
-    // Hack ends here
-
     private synchronized void createFunction(SqlInvokedFunction function)
     {
         checkFunctionLanguageSupported(function);
diff --git a/presto-native-sidecar-plugin/src/main/java/com/facebook/presto/sidecar/functionNamespace/NativeFunctionNamespaceManagerConfig.java b/presto-native-sidecar-plugin/src/main/java/com/facebook/presto/sidecar/functionNamespace/NativeFunctionNamespaceManagerConfig.java
new file mode 100644
index 0000000000000..38aeb75d1e5aa
--- /dev/null
+++ b/presto-native-sidecar-plugin/src/main/java/com/facebook/presto/sidecar/functionNamespace/NativeFunctionNamespaceManagerConfig.java
@@ -0,0 +1,52 @@
+/*
+ * Licensed 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 com.facebook.presto.sidecar.functionNamespace;
+
+import com.facebook.airlift.configuration.Config;
+import com.facebook.airlift.configuration.ConfigDescription;
+import com.facebook.airlift.units.Duration;
+
+import static java.util.concurrent.TimeUnit.MINUTES;
+
+public class NativeFunctionNamespaceManagerConfig
+{
+    private int sidecarNumRetries = 8;
+    private Duration sidecarRetryDelay = new Duration(1, MINUTES);
+
+    public int getSidecarNumRetries()
+    {
+        return sidecarNumRetries;
+    }
+
+    @Config("sidecar.num-retries")
+    @ConfigDescription("Max times to retry fetching sidecar node")
+    public NativeFunctionNamespaceManagerConfig setSidecarNumRetries(int sidecarNumRetries)
+    {
+        this.sidecarNumRetries = sidecarNumRetries;
+        return this;
+    }
+
+    public Duration getSidecarRetryDelay()
+    {
+        return sidecarRetryDelay;
+    }
+
+    @Config("sidecar.retry-delay")
+    @ConfigDescription("Cooldown period to retry when fetching sidecar node fails")
+    public NativeFunctionNamespaceManagerConfig setSidecarRetryDelay(Duration sidecarRetryDelay)
+    {
+        this.sidecarRetryDelay = sidecarRetryDelay;
+        return this;
+    }
+}
diff --git a/presto-native-sidecar-plugin/src/main/java/com/facebook/presto/sidecar/functionNamespace/NativeFunctionNamespaceManagerModule.java b/presto-native-sidecar-plugin/src/main/java/com/facebook/presto/sidecar/functionNamespace/NativeFunctionNamespaceManagerModule.java
index 50f36269456d3..bff0a8d91d820 100644
--- a/presto-native-sidecar-plugin/src/main/java/com/facebook/presto/sidecar/functionNamespace/NativeFunctionNamespaceManagerModule.java
+++ b/presto-native-sidecar-plugin/src/main/java/com/facebook/presto/sidecar/functionNamespace/NativeFunctionNamespaceManagerModule.java
@@ -37,7 +37,6 @@ public class NativeFunctionNamespaceManagerModule
         implements Module
 {
     private final String catalogName;
-
     private final NodeManager nodeManager;
     private final FunctionMetadataManager functionMetadataManager;
 
@@ -55,6 +54,7 @@ public void configure(Binder binder)
 
         configBinder(binder).bindConfig(SqlInvokedFunctionNamespaceManagerConfig.class);
         configBinder(binder).bindConfig(SqlFunctionLanguageConfig.class);
+        configBinder(binder).bindConfig(NativeFunctionNamespaceManagerConfig.class);
         binder.bind(FunctionDefinitionProvider.class).to(NativeFunctionDefinitionProvider.class).in(SINGLETON);
         binder.bind(new TypeLiteral>>>() {})
                 .toInstance(new JsonCodecFactory().mapJsonCodec(String.class, listJsonCodec(JsonBasedUdfFunctionMetadata.class)));
diff --git a/presto-native-sidecar-plugin/src/main/java/com/facebook/presto/sidecar/typemanager/NativeTypeManager.java b/presto-native-sidecar-plugin/src/main/java/com/facebook/presto/sidecar/typemanager/NativeTypeManager.java
index a1cd8cb0f00af..dd4122a04cc6d 100644
--- a/presto-native-sidecar-plugin/src/main/java/com/facebook/presto/sidecar/typemanager/NativeTypeManager.java
+++ b/presto-native-sidecar-plugin/src/main/java/com/facebook/presto/sidecar/typemanager/NativeTypeManager.java
@@ -20,7 +20,7 @@
 import com.facebook.presto.common.type.TypeManager;
 import com.facebook.presto.common.type.TypeSignature;
 import com.facebook.presto.common.type.TypeSignatureParameter;
-import com.facebook.presto.common.type.VarcharType;
+import com.facebook.presto.spi.type.UnknownTypeException;
 import com.google.common.cache.CacheBuilder;
 import com.google.common.cache.CacheLoader;
 import com.google.common.cache.LoadingCache;
@@ -66,7 +66,6 @@
 import static com.facebook.presto.common.type.StandardTypes.UUID;
 import static com.facebook.presto.common.type.StandardTypes.VARBINARY;
 import static com.facebook.presto.common.type.StandardTypes.VARCHAR;
-import static com.facebook.presto.common.type.VarcharType.createUnboundedVarcharType;
 import static com.google.common.base.Preconditions.checkArgument;
 import static com.google.common.base.Preconditions.checkState;
 import static com.google.common.base.Throwables.throwIfUnchecked;
@@ -108,7 +107,10 @@ public class NativeTypeManager
                     QDIGEST,
                     ROW,
                     TDIGEST,
-                    FunctionType.NAME);
+                    FunctionType.NAME,
+                    // todo: fix this hack: parametrized varchar isn't supported in native execution yet
+                    // Currently, native execution without sidecar enabled works with parameterized varchar, hence adding this to the list of supported types.
+                    VARCHAR);
 
     private final TypeManager typeManager;
     private final LoadingCache parametricTypeCache;
@@ -133,10 +135,6 @@ public NativeTypeManager(TypeManager typeManager)
     @Override
     public Type getType(TypeSignature typeSignature)
     {
-        // Todo: Fix this hack, native execution does not support parameterized varchar type signatures.
-        if (typeSignature.getBase().equals(VARCHAR)) {
-            typeSignature = createUnboundedVarcharType().getTypeSignature();
-        }
         Type type = types.get(typeSignature);
         if (type != null) {
             return type;
@@ -150,6 +148,18 @@ public Type getType(TypeSignature typeSignature)
         }
     }
 
+    @Override
+    public boolean hasType(TypeSignature typeSignature)
+    {
+        try {
+            getType(typeSignature);
+            return true;
+        }
+        catch (UnknownTypeException e) {
+            return false;
+        }
+    }
+
     @Override
     public Collection getParametricTypes()
     {
@@ -177,10 +187,6 @@ public boolean canCoerce(Type actualType, Type expectedType)
     private void addAllTypes(List typesList, List parametricTypesList)
     {
         typesList.forEach(this::addType);
-        // todo: Fix this hack
-        // Native engine does not support parameterized varchar, and varchar isn't in the lists of types returned from the engine
-        addType(VarcharType.VARCHAR);
-
         parametricTypesList.forEach(this::addParametricType);
     }
 
diff --git a/presto-native-sidecar-plugin/src/test/java/com/facebook/presto/sidecar/NativeSidecarPluginQueryRunnerUtils.java b/presto-native-sidecar-plugin/src/test/java/com/facebook/presto/sidecar/NativeSidecarPluginQueryRunnerUtils.java
index 776d4920e2f16..c8c7e1123f974 100644
--- a/presto-native-sidecar-plugin/src/test/java/com/facebook/presto/sidecar/NativeSidecarPluginQueryRunnerUtils.java
+++ b/presto-native-sidecar-plugin/src/test/java/com/facebook/presto/sidecar/NativeSidecarPluginQueryRunnerUtils.java
@@ -13,6 +13,7 @@
  */
 package com.facebook.presto.sidecar;
 
+import com.facebook.presto.scalar.sql.NativeSqlInvokedFunctionsPlugin;
 import com.facebook.presto.sidecar.functionNamespace.NativeFunctionNamespaceManagerFactory;
 import com.facebook.presto.sidecar.sessionpropertyproviders.NativeSystemSessionPropertyProviderFactory;
 import com.facebook.presto.sidecar.typemanager.NativeTypeManagerFactory;
@@ -37,5 +38,6 @@ public static void setupNativeSidecarPlugin(QueryRunner queryRunner)
                         "function-implementation-type", "CPP"));
         queryRunner.loadTypeManager(NativeTypeManagerFactory.NAME);
         queryRunner.loadPlanCheckerProviderManager("native", ImmutableMap.of());
+        queryRunner.installPlugin(new NativeSqlInvokedFunctionsPlugin());
     }
 }
diff --git a/presto-native-sidecar-plugin/src/test/java/com/facebook/presto/sidecar/TestNativeSidecarPlugin.java b/presto-native-sidecar-plugin/src/test/java/com/facebook/presto/sidecar/TestNativeSidecarPlugin.java
index 05945f225456e..d3f584222c3b9 100644
--- a/presto-native-sidecar-plugin/src/test/java/com/facebook/presto/sidecar/TestNativeSidecarPlugin.java
+++ b/presto-native-sidecar-plugin/src/test/java/com/facebook/presto/sidecar/TestNativeSidecarPlugin.java
@@ -14,7 +14,10 @@
 package com.facebook.presto.sidecar;
 
 import com.facebook.airlift.units.DataSize;
+import com.facebook.presto.Session;
 import com.facebook.presto.nativeworker.PrestoNativeQueryRunnerUtils;
+import com.facebook.presto.scalar.sql.NativeSqlInvokedFunctionsPlugin;
+import com.facebook.presto.scalar.sql.SqlInvokedFunctionsPlugin;
 import com.facebook.presto.sidecar.functionNamespace.FunctionDefinitionProvider;
 import com.facebook.presto.sidecar.functionNamespace.NativeFunctionDefinitionProvider;
 import com.facebook.presto.sidecar.functionNamespace.NativeFunctionNamespaceManager;
@@ -44,8 +47,12 @@
 import java.util.stream.Collectors;
 
 import static com.facebook.airlift.units.DataSize.Unit.MEGABYTE;
+import static com.facebook.presto.SystemSessionProperties.INLINE_SQL_FUNCTIONS;
+import static com.facebook.presto.SystemSessionProperties.KEY_BASED_SAMPLING_ENABLED;
+import static com.facebook.presto.SystemSessionProperties.REMOVE_MAP_CAST;
 import static com.facebook.presto.common.Utils.checkArgument;
 import static com.facebook.presto.common.type.BigintType.BIGINT;
+import static com.facebook.presto.nativeworker.NativeQueryRunnerUtils.createCustomer;
 import static com.facebook.presto.nativeworker.NativeQueryRunnerUtils.createLineitem;
 import static com.facebook.presto.nativeworker.NativeQueryRunnerUtils.createNation;
 import static com.facebook.presto.nativeworker.NativeQueryRunnerUtils.createOrders;
@@ -63,6 +70,7 @@ public class TestNativeSidecarPlugin
     private static final String REGEX_FUNCTION_NAMESPACE = "native.default.*";
     private static final String REGEX_SESSION_NAMESPACE = "Native Execution only.*";
     private static final long SIDECAR_HTTP_CLIENT_MAX_CONTENT_SIZE_MB = 128;
+    private static final int INLINED_SQL_FUNCTIONS_COUNT = 7;
 
     @Override
     protected void createTables()
@@ -73,6 +81,7 @@ protected void createTables()
         createOrders(queryRunner);
         createOrdersEx(queryRunner);
         createRegion(queryRunner);
+        createCustomer(queryRunner);
     }
 
     @Override
@@ -91,9 +100,11 @@ protected QueryRunner createQueryRunner()
     protected QueryRunner createExpectedQueryRunner()
             throws Exception
     {
-        return PrestoNativeQueryRunnerUtils.javaHiveQueryRunnerBuilder()
+        QueryRunner queryRunner = PrestoNativeQueryRunnerUtils.javaHiveQueryRunnerBuilder()
                 .setAddStorageFormatToPath(true)
                 .build();
+        queryRunner.installPlugin(new SqlInvokedFunctionsPlugin());
+        return queryRunner;
     }
 
     public static void setupNativeSidecarPlugin(QueryRunner queryRunner)
@@ -111,6 +122,7 @@ public static void setupNativeSidecarPlugin(QueryRunner queryRunner)
                         "sidecar.http-client.max-content-length", SIDECAR_HTTP_CLIENT_MAX_CONTENT_SIZE_MB + "MB"));
         queryRunner.loadTypeManager(NativeTypeManagerFactory.NAME);
         queryRunner.loadPlanCheckerProviderManager("native", ImmutableMap.of());
+        queryRunner.installPlugin(new NativeSqlInvokedFunctionsPlugin());
     }
 
     @Test
@@ -155,12 +167,13 @@ public void testSetNativeWorkerSessionProperty()
                 "MaterializedResult{rows=[[true]], " +
                         "types=[boolean], " +
                         "setSessionProperties={driver_cpu_time_slice_limit_ms=500}, " +
-                        "resetSessionProperties=[], updateType=SET SESSION}");
+                        "resetSessionProperties=[], updateType=SET SESSION, clearTransactionId=false}");
     }
 
     @Test
     public void testShowFunctions()
     {
+        int inlinedSQLFunctionsCount = 0;
         @Language("SQL") String sql = "SHOW FUNCTIONS";
         MaterializedResult actualResult = computeActual(sql);
         List actualRows = actualResult.getMaterializedRows();
@@ -174,11 +187,17 @@ public void testShowFunctions()
 
             // function namespace should be present.
             String fullFunctionName = row.get(5).toString();
-            if (Pattern.matches(REGEX_FUNCTION_NAMESPACE, fullFunctionName)) {
-                continue;
+            if (!Pattern.matches(REGEX_FUNCTION_NAMESPACE, fullFunctionName)) {
+                // If no namespace match found, check if it's an inlined SQL Invoked function.
+                String language = row.get(9).toString();
+                if (language.equalsIgnoreCase("SQL")) {
+                    inlinedSQLFunctionsCount++;
+                    continue;
+                }
+                fail(format("No namespace match found for row: %s", row));
             }
-            fail(format("No namespace match found for row: %s", row));
         }
+        assertEquals(inlinedSQLFunctionsCount, INLINED_SQL_FUNCTIONS_COUNT);
     }
 
     @Test
@@ -201,6 +220,7 @@ public void testGeneralQueries()
                 "date_trunc('second', from_unixtime(orderkey, '+00:00')) FROM orders");
         assertQuery("SELECT mod(orderkey, linenumber) FROM lineitem");
         assertQueryFails("SELECT IF(true, 0/0, 1)", "[\\s\\S]*/ by zero native.default.fail[\\s\\S]*");
+        assertQuery("select CASE WHEN true THEN 'Yes' ELSE 'No' END");
     }
 
     @Test
@@ -269,6 +289,70 @@ public void testLambdaFunctions()
                         " Exception 2: line 1:31: Expected a lambda that takes ([12])" + Pattern.quote(" argument(s) but got 3\n"));
     }
 
+    @Test
+    public void testArraySortByKeyFunction()
+    {
+        // Basic string sorting by length
+        assertQuerySucceeds("SELECT array_sort(ARRAY['pear', 'apple', 'banana', 'kiwi'], x -> length(x))");
+        assertQuerySucceeds("SELECT array_sort(ARRAY['pear', 'apple', 'banana', 'kiwi'], x -> substr(x, length(x), 1))");
+        // Sorting with nulls
+        assertQuerySucceeds("SELECT array_sort(ARRAY['apple', NULL, 'banana', NULL], x -> length(x))");
+        assertQuerySucceeds("SELECT array_sort(ARRAY['apple', 'banana', 'pear'], x -> IF(x = 'banana', NULL, length(x)))");
+        assertQuerySucceeds("SELECT array_sort(ARRAY['apple', NULL, 'banana', 'pear', NULL], x -> length(x))");
+        // Special double values
+        assertQuerySucceeds("SELECT array_sort(ARRAY[CAST(0.0 AS DOUBLE), CAST('NaN' AS DOUBLE), CAST('Infinity' AS DOUBLE), CAST('-Infinity' AS DOUBLE)], x -> x)");
+        // Numeric keys
+        assertQuerySucceeds("SELECT array_sort(ARRAY[5, 20, 3, 9, 100], x -> x)");
+        assertQuerySucceeds("SELECT array_sort(ARRAY[CAST(5000000000 AS BIGINT), CAST(20000000000 AS BIGINT), CAST(3000000000 AS BIGINT), CAST(9000000000 AS BIGINT), CAST(100000000000 AS BIGINT)], x -> x)");
+        assertQuerySucceeds("SELECT array_sort(ARRAY[CAST(5.5 AS DOUBLE), CAST(20.1 AS DOUBLE), CAST(3.9 AS DOUBLE), CAST(9.0 AS DOUBLE), CAST(100.0 AS DOUBLE)], x -> x)");
+        assertQuerySucceeds("SELECT array_sort(ARRAY[5, 20, 3, 9, 100], x -> x % 10)");
+        // Boolean keys
+        assertQuerySucceeds("SELECT array_sort(ARRAY[true, false, true, false], x -> x)");
+        assertQuerySucceeds("SELECT array_sort(ARRAY[true, false, true, false], x -> NOT x)");
+        // Complex types
+        assertQuerySucceeds("SELECT array_sort(ARRAY[ARRAY[1, 2, 3], ARRAY[4, 5], ARRAY[6, 7, 8, 9]], x -> cardinality(x))");
+        assertQuerySucceeds("SELECT array_sort(ARRAY[ROW('a', 3), ROW('b', 1), ROW('c', 2)], x -> x[2])");
+        // Edge cases
+        assertQuerySucceeds("SELECT array_sort(ARRAY[], x -> x)");
+        assertQuerySucceeds("SELECT array_sort(ARRAY[5], x -> x)");
+        assertQuerySucceeds("SELECT array_sort(ARRAY[NULL, NULL, NULL], x -> x)");
+        // Type coercion
+        assertQuerySucceeds("SELECT array_sort(ARRAY[5, 20, 3, 9, 100], x -> x + CAST(0.5 AS DOUBLE))");
+        assertQuerySucceeds("SELECT array_sort(ARRAY[5, 20, 3, 9, 100], x -> x * CAST(1000000000 AS BIGINT))");
+    }
+
+    @Test
+    public void testArraySortDescByKeyFunction()
+    {
+        // Basic string sorting by length in descending order
+        assertQuerySucceeds("SELECT array_sort_desc(ARRAY['pear', 'apple', 'banana', 'kiwi'], x -> length(x))");
+        assertQuerySucceeds("SELECT array_sort_desc(ARRAY['pear', 'apple', 'banana', 'kiwi'], x -> substr(x, length(x), 1))");
+        // Sorting with nulls
+        assertQuerySucceeds("SELECT array_sort_desc(ARRAY['apple', NULL, 'banana', NULL], x -> length(x))");
+        assertQuerySucceeds("SELECT array_sort_desc(ARRAY['apple', 'banana', 'pear'], x -> IF(x = 'banana', NULL, length(x)))");
+        assertQuerySucceeds("SELECT array_sort_desc(ARRAY['apple', NULL, 'banana', 'pear', NULL], x -> length(x))");
+        // Special double values
+        assertQuerySucceeds("SELECT array_sort_desc(ARRAY[CAST(0.0 AS DOUBLE), CAST('NaN' AS DOUBLE), CAST('Infinity' AS DOUBLE), CAST('-Infinity' AS DOUBLE)], x -> x)");
+        // Numeric keys
+        assertQuerySucceeds("SELECT array_sort_desc(ARRAY[5, 20, 3, 9, 100], x -> x)");
+        assertQuerySucceeds("SELECT array_sort_desc(ARRAY[CAST(5000000000 AS BIGINT), CAST(20000000000 AS BIGINT), CAST(3000000000 AS BIGINT), CAST(9000000000 AS BIGINT), CAST(100000000000 AS BIGINT)], x -> x)");
+        assertQuerySucceeds("SELECT array_sort_desc(ARRAY[CAST(5.5 AS DOUBLE), CAST(20.1 AS DOUBLE), CAST(3.9 AS DOUBLE), CAST(9.0 AS DOUBLE), CAST(100.0 AS DOUBLE)], x -> x)");
+        assertQuerySucceeds("SELECT array_sort_desc(ARRAY[5, 20, 3, 9, 100], x -> x % 10)");
+        // Boolean keys
+        assertQuerySucceeds("SELECT array_sort_desc(ARRAY[true, false, true, false], x -> x)");
+        assertQuerySucceeds("SELECT array_sort_desc(ARRAY[true, false, true, false], x -> NOT x)");
+        // Complex types
+        assertQuerySucceeds("SELECT array_sort_desc(ARRAY[ARRAY[1, 2, 3], ARRAY[4, 5], ARRAY[6, 7, 8, 9]], x -> cardinality(x))");
+        assertQuerySucceeds("SELECT array_sort_desc(ARRAY[ROW('a', 3), ROW('b', 1), ROW('c', 2)], x -> x[2])");
+        // Edge cases
+        assertQuerySucceeds("SELECT array_sort_desc(ARRAY[], x -> x)");
+        assertQuerySucceeds("SELECT array_sort_desc(ARRAY[5], x -> x)");
+        assertQuerySucceeds("SELECT array_sort_desc(ARRAY[NULL, NULL, NULL], x -> x)");
+        // Type coercion
+        assertQuerySucceeds("SELECT array_sort_desc(ARRAY[5, 20, 3, 9, 100], x -> x + CAST(0.5 AS DOUBLE))");
+        assertQuerySucceeds("SELECT array_sort_desc(ARRAY[5, 20, 3, 9, 100], x -> x * CAST(1000000000 AS BIGINT))");
+    }
+
     @Test
     public void testApproxPercentile()
     {
@@ -315,11 +399,30 @@ public void testApproxPercentile()
         }
     }
 
+    @Test
+    public void testMapSubset()
+    {
+        assertQuery("select m[1], m[3] from (select map_subset(map(array[1,2,3,4], array['a', 'b', 'c', 'd']), array[1,3,10]) m)", "select 'a', 'c'");
+        assertQuery("select m['p'], m['r'] from (select map_subset(map(array['p', 'q', 'r', 's'], array['a', 'b', 'c', 'd']), array['p', 'r', 'z']) m)", "select 'a', 'c'");
+        assertQuery("select m[true], m[false] from (select map_subset(map(array[false, true], array['a', 'z']), array[true, false]) m)", "select 'z', 'a'");
+        assertQuery("select m[DATE '2015-01-01'], m[DATE '2015-01-13'] from (select map_subset(" +
+                "map(array[DATE '2015-01-01', DATE '2015-02-13', DATE '2015-01-13', DATE '2015-05-15'], array['a', 'b', 'c', 'd']), " +
+                "array[DATE '2015-01-01', DATE '2015-01-13', DATE '2015-06-15']) m)", "select 'a', 'c'");
+        assertQuery("select m[TIMESTAMP '2021-01-02 09:04:05.321'] from (select map_subset(" +
+                "map(array[TIMESTAMP '2021-01-02 09:04:05.321', TIMESTAMP '2022-12-22 10:07:08.456'], array['a', 'b']), " +
+                "array[TIMESTAMP '2021-01-02 09:04:05.321', TIMESTAMP '2022-12-22 10:07:09.246']) m)", "select 'a'");
+    }
+
     @Test
     public void testInformationSchemaTables()
     {
         assertQuery("select lower(table_name) from information_schema.tables "
-                        + "where table_name = 'lineitem' or table_name = 'LINEITEM' ");
+                + "where table_name = 'lineitem' or table_name = 'LINEITEM' ");
+        assertQuery("SELECT table_name, CASE WHEN abs(ordinal_position) > 3 THEN 'high' WHEN abs(ordinal_position) > 1 THEN 'medium' ELSE 'low' END as position_category, COUNT(*) \n" +
+                "FROM information_schema.columns " +
+                "WHERE table_catalog = 'hive' AND table_name IN ('nation', 'region', 'lineitem', 'orders') " +
+                "GROUP BY table_name, CASE WHEN abs(ordinal_position) > 3 THEN 'high' WHEN abs(ordinal_position) > 1 THEN 'medium' ELSE 'low' END " +
+                "ORDER BY table_name, position_category");
     }
 
     @Test
@@ -393,6 +496,130 @@ public void testGeometryQueries()
         assertQuery("SELECT " +
                 "ST_DISTANCE(ST_POINT(a.nationkey, a.regionkey), ST_POINT(b.nationkey, b.regionkey)) " +
                 "FROM nation a JOIN nation b ON a.nationkey < b.nationkey");
+        assertQuery(
+                "WITH regions(name, geom) AS (VALUES" +
+                        "        ('A', ST_GeometryFromText('POLYGON ((0 0, 0 5, 5 5, 5 0, 0 0))'))," +
+                        "        ('B', ST_GeometryFromText('POLYGON ((5 0, 5 5, 10 5, 10 0, 5 0))')))," +
+                        "points(id, geom) AS (VALUES" +
+                        "        ('P1', ST_Point(1, 1))," +
+                        "        ('P2', ST_Point(6, 1))," +
+                        "        ('P3', ST_Point(8, 4)))" +
+                        "SELECT p.id, r.name FROM points p LEFT JOIN regions r ON ST_Within(p.geom, r.geom)");
+    }
+
+    @Test
+    public void testRemoveMapCast()
+    {
+        Session enableOptimization = Session.builder(getSession())
+                .setSystemProperty(REMOVE_MAP_CAST, "true")
+                .build();
+        assertQuery(enableOptimization, "select feature[key] from (values (map(array[cast(1 as integer), 2, 3, 4], array[0.3, 0.5, 0.9, 0.1]), cast(2 as bigint)), (map(array[cast(1 as integer), 2, 3, 4], array[0.3, 0.5, 0.9, 0.1]), 4)) t(feature,  key)",
+                "values 0.5, 0.1");
+        assertQuery(enableOptimization, "select element_at(feature, key) from (values (map(array[cast(1 as integer), 2, 3, 4], array[0.3, 0.5, 0.9, 0.1]), cast(2 as bigint)), (map(array[cast(1 as integer), 2, 3, 4], array[0.3, 0.5, 0.9, 0.1]), 4)) t(feature,  key)",
+                "values 0.5, 0.1");
+        assertQuery(enableOptimization, "select element_at(feature, key) from (values (map(array[cast(1 as integer), 2, 3, 4], array[0.3, 0.5, 0.9, 0.1]), cast(2 as bigint)), (map(array[cast(1 as integer), 2, 3, 4], array[0.3, 0.5, 0.9, 0.1]), 400000000000)) t(feature, key)",
+                "values 0.5, null");
+        assertQuery(enableOptimization, "select feature[key] from (values (map(array[cast(1 as varchar), '2', '3', '4'], array[0.3, 0.5, 0.9, 0.1]), cast('2' as varchar)), (map(array[cast(1 as varchar), '2', '3', '4'], array[0.3, 0.5, 0.9, 0.1]), '4')) t(feature,  key)",
+                "values 0.5, 0.1");
+    }
+
+    @Test
+    public void testOverriddenInlinedSqlInvokedFunctions()
+    {
+        // String functions
+        assertQuery("SELECT trail(comment, cast(nationkey as integer)) FROM nation");
+        assertQuery("SELECT name, comment, replace_first(comment, 'iron', 'gold') from nation");
+
+        // Array functions
+        assertQuery("SELECT array_intersect(ARRAY['apple', 'banana', 'cherry'], ARRAY['apple', 'mango', 'fig'])");
+        assertQuery("SELECT array_frequency(split(comment, '')) from nation");
+        assertQuery("SELECT array_duplicates(ARRAY[regionkey]), array_duplicates(ARRAY[comment]) from nation");
+        assertQuery("SELECT array_has_duplicates(ARRAY[custkey]) from orders");
+        assertQuery("SELECT array_max_by(ARRAY[comment], x -> length(x)) from orders");
+        assertQuery("SELECT array_min_by(ARRAY[ROW('USA', 1), ROW('INDIA', 2), ROW('UK', 3)], x -> x[2])");
+        assertQuery("SELECT array_sort_desc(map_keys(map_union(quantity_by_linenumber))) FROM orders_ex");
+        assertQuery("SELECT remove_nulls(ARRAY[CAST(regionkey AS VARCHAR), comment, NULL]) from nation");
+        assertQuery("SELECT array_top_n(ARRAY[CAST(nationkey AS VARCHAR)], 3) from nation");
+        assertQuerySucceeds("SELECT array_sort_desc(quantities, x -> abs(x)) FROM orders_ex");
+
+        // Map functions
+        assertQuery("SELECT map_normalize(MAP(ARRAY['a', 'b', 'c'], ARRAY[1, 4, 5]))");
+        assertQuery("SELECT map_normalize(MAP(ARRAY['a', 'b', 'c'], ARRAY[1, 0, -1]))");
+        assertQuery("SELECT name, map_normalize(MAP(ARRAY['regionkey', 'length'], ARRAY[regionkey, length(comment)])) from nation");
+        assertQuery("SELECT name, map_remove_null_values(map(ARRAY['region', 'comment', 'nullable'], " +
+                "ARRAY[CAST(regionkey AS VARCHAR), comment, NULL])) from nation");
+        assertQuery("SELECT name, map_key_exists(map(ARRAY['nation', 'comment'], ARRAY[CAST(nationkey AS VARCHAR), comment]), 'comment') from nation");
+        assertQuery("SELECT map_keys_by_top_n_values(MAP(ARRAY[orderkey], ARRAY[custkey]), 2) from orders");
+        assertQuery("SELECT map_top_n(MAP(ARRAY[CAST(nationkey AS VARCHAR)], ARRAY[comment]), 3) from nation");
+        assertQuery("SELECT map_top_n_keys(MAP(ARRAY[orderkey], ARRAY[custkey]), 3) from orders");
+        assertQuery("SELECT map_top_n_values(MAP(ARRAY[orderkey], ARRAY[custkey]), 3) from orders");
+        assertQuery("SELECT all_keys_match(MAP(ARRAY[comment], ARRAY[custkey]), k -> length(k) > 5) from orders");
+        assertQuery("SELECT any_keys_match(MAP(ARRAY[comment], ARRAY[custkey]), k -> starts_with(k, 'abc')) from orders");
+        assertQuery("SELECT any_values_match(MAP(ARRAY[orderkey], ARRAY[totalprice]), k -> abs(k) > 20) from orders");
+        assertQuery("SELECT no_values_match(MAP(ARRAY[orderkey], ARRAY[comment]), k -> length(k) > 2) from orders");
+        assertQuery("SELECT no_keys_match(MAP(ARRAY[comment], ARRAY[custkey]), k -> ends_with(k, 'a')) from orders");
+    }
+
+    @Test
+    public void testNonOverriddenInlinedSqlInvokedFunctionsWhenConfigEnabled()
+    {
+        // Array functions
+        assertQuery("SELECT array_split_into_chunks(split(comment, ''), 2) from nation");
+        assertQuery("SELECT array_least_frequent(quantities) from orders_ex");
+        assertQuery("SELECT array_least_frequent(split(comment, ''), 5) from nation");
+        assertQuerySucceeds("SELECT array_top_n(ARRAY[orderkey], 25, (x, y) -> if (x < y, cast(1 as bigint), if (x > y, cast(-1 as bigint), cast(0 as bigint)))) from orders");
+
+        // Map functions
+        assertQuerySucceeds("SELECT map_top_n_values(MAP(ARRAY[comment], ARRAY[nationkey]), 2, (x, y) -> if (x < y, cast(1 as bigint), if (x > y, cast(-1 as bigint), cast(0 as bigint)))) from nation");
+        assertQuerySucceeds("SELECT map_top_n_keys(MAP(ARRAY[regionkey], ARRAY[nationkey]), 5, (x, y) -> if (x < y, cast(1 as bigint), if (x > y, cast(-1 as bigint), cast(0 as bigint)))) from nation");
+
+        Session sessionWithKeyBasedSampling = Session.builder(getSession())
+                .setSystemProperty(KEY_BASED_SAMPLING_ENABLED, "true")
+                .build();
+
+        @Language("SQL") String query = "select count(1) FROM lineitem l left JOIN orders o ON l.orderkey = o.orderkey JOIN customer c ON o.custkey = c.custkey";
+
+        assertQuery(query, "select cast(60175 as bigint)");
+        assertQuery(sessionWithKeyBasedSampling, query, "select cast(16185 as bigint)");
+    }
+
+    @Test
+    public void testNonOverriddenInlinedSqlInvokedFunctionsWhenConfigDisabled()
+    {
+        // When inline_sql_functions is set to false, the below queries should fail as the implementations don't exist on the native worker
+        Session session = Session.builder(getSession())
+                .setSystemProperty(KEY_BASED_SAMPLING_ENABLED, "true")
+                .setSystemProperty(INLINE_SQL_FUNCTIONS, "false")
+                .build();
+
+        // Array functions
+        assertQueryFails(session,
+                "SELECT array_split_into_chunks(split(comment, ''), 2) from nation",
+                ".*Scalar function name not registered: native.default.array_split_into_chunks.*");
+        assertQueryFails(session,
+                "SELECT array_least_frequent(quantities) from orders_ex",
+                ".*Scalar function name not registered: native.default.array_least_frequent.*");
+        assertQueryFails(session,
+                "SELECT array_least_frequent(split(comment, ''), 2) from nation",
+                ".*Scalar function name not registered: native.default.array_least_frequent.*");
+        assertQueryFails(session,
+                "SELECT array_top_n(ARRAY[orderkey], 25, (x, y) -> if (x < y, cast(1 as bigint), if (x > y, cast(-1 as bigint), cast(0 as bigint)))) from orders",
+                " Scalar function native\\.default\\.array_top_n not registered with arguments.*",
+                true);
+
+        // Map functions
+        assertQueryFails(session,
+                "SELECT map_top_n_values(MAP(ARRAY[comment], ARRAY[nationkey]), 2, (x, y) -> if (x < y, cast(1 as bigint), if (x > y, cast(-1 as bigint), cast(0 as bigint)))) from nation",
+                ".*Scalar function native\\.default\\.map_top_n_values not registered with arguments.*",
+                true);
+        assertQueryFails(session,
+                "SELECT map_top_n_keys(MAP(ARRAY[regionkey], ARRAY[nationkey]), 5, (x, y) -> if (x < y, cast(1 as bigint), if (x > y, cast(-1 as bigint), cast(0 as bigint)))) from nation",
+                ".*Scalar function native\\.default\\.map_top_n_keys not registered with arguments.*",
+                true);
+
+        assertQueryFails(session,
+                "select count(1) FROM lineitem l left JOIN orders o ON l.orderkey = o.orderkey JOIN customer c ON o.custkey = c.custkey",
+                ".*Scalar function name not registered: native.default.key_sampling_percent.*");
     }
 
     private String generateRandomTableName()
diff --git a/presto-native-sidecar-plugin/src/test/java/com/facebook/presto/sidecar/TestNativeSidecarPluginWithoutLoadingFunctionalities.java b/presto-native-sidecar-plugin/src/test/java/com/facebook/presto/sidecar/TestNativeSidecarPluginWithoutLoadingFunctionalities.java
index 01e786a4a0f66..a0997d9ab365c 100644
--- a/presto-native-sidecar-plugin/src/test/java/com/facebook/presto/sidecar/TestNativeSidecarPluginWithoutLoadingFunctionalities.java
+++ b/presto-native-sidecar-plugin/src/test/java/com/facebook/presto/sidecar/TestNativeSidecarPluginWithoutLoadingFunctionalities.java
@@ -75,9 +75,5 @@ public void testBasicQueries()
         assertQuery("select count(*) from nation");
         assertQuery("select count(abs(orderkey) between 1 and 60000) from orders group by orderkey");
         assertQuery("SELECT count(orderkey) FROM orders WHERE orderkey < 0 GROUP BY GROUPING SETS (())");
-        // This query will work when sidecar is enabled, should fail without it.
-        assertQueryFails(
-                "select array_sort(array[row('apples', 23), row('bananas', 12), row('grapes', 44)], x -> x[2])",
-                "line 1:84: Expected a lambda that takes 2 argument\\(s\\) but got 1");
     }
 }
diff --git a/presto-native-sidecar-plugin/src/test/java/com/facebook/presto/sidecar/functionNamespace/TestNativeFunctionNamespaceManagerConfig.java b/presto-native-sidecar-plugin/src/test/java/com/facebook/presto/sidecar/functionNamespace/TestNativeFunctionNamespaceManagerConfig.java
new file mode 100644
index 0000000000000..feee5f6ce8c9c
--- /dev/null
+++ b/presto-native-sidecar-plugin/src/test/java/com/facebook/presto/sidecar/functionNamespace/TestNativeFunctionNamespaceManagerConfig.java
@@ -0,0 +1,51 @@
+/*
+ * Licensed 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 com.facebook.presto.sidecar.functionNamespace;
+
+import com.facebook.airlift.units.Duration;
+import com.google.common.collect.ImmutableMap;
+import org.testng.annotations.Test;
+
+import java.util.Map;
+
+import static com.facebook.airlift.configuration.testing.ConfigAssertions.assertFullMapping;
+import static com.facebook.airlift.configuration.testing.ConfigAssertions.assertRecordedDefaults;
+import static com.facebook.airlift.configuration.testing.ConfigAssertions.recordDefaults;
+import static java.util.concurrent.TimeUnit.MINUTES;
+
+public class TestNativeFunctionNamespaceManagerConfig
+{
+    @Test
+    public void testDefaults()
+    {
+        assertRecordedDefaults(recordDefaults(NativeFunctionNamespaceManagerConfig.class)
+                .setSidecarNumRetries(8)
+                .setSidecarRetryDelay(new Duration(1, MINUTES)));
+    }
+
+    @Test
+    public void testExplicitPropertyMappings()
+    {
+        Map properties = new ImmutableMap.Builder()
+                .put("sidecar.num-retries", "15")
+                .put("sidecar.retry-delay", "5m")
+                .build();
+
+        NativeFunctionNamespaceManagerConfig expected = new NativeFunctionNamespaceManagerConfig()
+                .setSidecarNumRetries(15)
+                .setSidecarRetryDelay(new Duration(5, MINUTES));
+
+        assertFullMapping(properties, expected);
+    }
+}
diff --git a/presto-native-sql-invoked-functions-plugin/pom.xml b/presto-native-sql-invoked-functions-plugin/pom.xml
new file mode 100644
index 0000000000000..7d837a28a8bfb
--- /dev/null
+++ b/presto-native-sql-invoked-functions-plugin/pom.xml
@@ -0,0 +1,29 @@
+
+    4.0.0
+    
+        com.facebook.presto
+        presto-root
+        0.295-SNAPSHOT
+    
+
+    presto-native-sql-invoked-functions-plugin
+    Presto Native - Sql invoked functions plugin
+    presto-plugin
+
+    
+        ${project.parent.basedir}
+    
+
+    
+        
+            com.facebook.presto
+            presto-spi
+            provided
+        
+        
+            com.google.guava
+            guava
+        
+    
+
diff --git a/presto-native-sql-invoked-functions-plugin/src/main/java/com/facebook/presto/scalar/sql/NativeArraySqlFunctions.java b/presto-native-sql-invoked-functions-plugin/src/main/java/com/facebook/presto/scalar/sql/NativeArraySqlFunctions.java
new file mode 100644
index 0000000000000..841883d99ae8f
--- /dev/null
+++ b/presto-native-sql-invoked-functions-plugin/src/main/java/com/facebook/presto/scalar/sql/NativeArraySqlFunctions.java
@@ -0,0 +1,74 @@
+/*
+ * Licensed 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 com.facebook.presto.scalar.sql;
+
+import com.facebook.presto.spi.function.Description;
+import com.facebook.presto.spi.function.SqlInvokedScalarFunction;
+import com.facebook.presto.spi.function.SqlParameter;
+import com.facebook.presto.spi.function.SqlParameters;
+import com.facebook.presto.spi.function.SqlType;
+import com.facebook.presto.spi.function.TypeParameter;
+
+public class NativeArraySqlFunctions
+{
+    private NativeArraySqlFunctions() {}
+
+    @SqlInvokedScalarFunction(value = "array_split_into_chunks", deterministic = true, calledOnNullInput = false)
+    @Description("Returns an array of arrays splitting input array into chunks of given length. " +
+            "If array is not evenly divisible it will split into as many possible chunks and " +
+            "return the left over elements for the last array. Returns null for null inputs, but not elements.")
+    @TypeParameter("T")
+    @SqlParameters({@SqlParameter(name = "input", type = "array(T)"), @SqlParameter(name = "sz", type = "int")})
+    @SqlType("array(array(T))")
+    public static String arraySplitIntoChunks()
+    {
+        return "RETURN IF(sz <= 0, " +
+                "fail('Invalid slice size: ' || cast(sz as varchar) || '. Size must be greater than zero.'), " +
+                "IF(cardinality(input) / sz > 10000, " +
+                "fail('Cannot split array of size: ' || cast(cardinality(input) as varchar) || ' into more than 10000 parts.'), " +
+                "transform(" +
+                "sequence(1, cardinality(input), sz), " +
+                "x -> slice(input, x, sz))))";
+    }
+
+    @SqlInvokedScalarFunction(value = "array_least_frequent", deterministic = true, calledOnNullInput = true)
+    @Description("Determines the least frequent element in the array. If there are multiple elements, the function returns the smallest element")
+    @TypeParameter("T")
+    @SqlParameter(name = "input", type = "array(T)")
+    @SqlType("array")
+    public static String arrayLeastFrequent()
+    {
+        return "RETURN IF(COALESCE(CARDINALITY(REMOVE_NULLS(input)), 0) = 0, NULL, TRANSFORM(SLICE(ARRAY_SORT(TRANSFORM(MAP_ENTRIES(ARRAY_FREQUENCY(REMOVE_NULLS(input))), x -> ROW(x[2], x[1]))), 1, 1), x -> x[2]))";
+    }
+
+    @SqlInvokedScalarFunction(value = "array_least_frequent", deterministic = true, calledOnNullInput = true)
+    @Description("Determines the n least frequent element in the array in the ascending order of the elements.")
+    @TypeParameter("T")
+    @SqlParameters({@SqlParameter(name = "input", type = "array(T)"), @SqlParameter(name = "n", type = "bigint")})
+    @SqlType("array")
+    public static String arrayNLeastFrequent()
+    {
+        return "RETURN IF(n < 0, fail('n must be greater than or equal to 0'), IF(COALESCE(CARDINALITY(REMOVE_NULLS(input)), 0) = 0, NULL, TRANSFORM(SLICE(ARRAY_SORT(TRANSFORM(MAP_ENTRIES(ARRAY_FREQUENCY(REMOVE_NULLS(input))), x -> ROW(x[2], x[1]))), 1, n), x -> x[2])))";
+    }
+
+    @SqlInvokedScalarFunction(value = "array_top_n", deterministic = true, calledOnNullInput = true)
+    @Description("Returns the top N values of the given map sorted using the provided lambda comparator.")
+    @TypeParameter("T")
+    @SqlParameters({@SqlParameter(name = "input", type = "array(T)"), @SqlParameter(name = "n", type = "int"), @SqlParameter(name = "f", type = "function(T, T, bigint)")})
+    @SqlType("array")
+    public static String arrayTopNComparator()
+    {
+        return "RETURN IF(n < 0, fail('Parameter n: ' || cast(n as varchar) || ' to ARRAY_TOP_N is negative'), SLICE(REVERSE(ARRAY_SORT(input, f)), 1, n))";
+    }
+}
diff --git a/presto-native-sql-invoked-functions-plugin/src/main/java/com/facebook/presto/scalar/sql/NativeMapSqlFunctions.java b/presto-native-sql-invoked-functions-plugin/src/main/java/com/facebook/presto/scalar/sql/NativeMapSqlFunctions.java
new file mode 100644
index 0000000000000..9eccc84d6d8c8
--- /dev/null
+++ b/presto-native-sql-invoked-functions-plugin/src/main/java/com/facebook/presto/scalar/sql/NativeMapSqlFunctions.java
@@ -0,0 +1,48 @@
+/*
+ * Licensed 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 com.facebook.presto.scalar.sql;
+
+import com.facebook.presto.spi.function.Description;
+import com.facebook.presto.spi.function.SqlInvokedScalarFunction;
+import com.facebook.presto.spi.function.SqlParameter;
+import com.facebook.presto.spi.function.SqlParameters;
+import com.facebook.presto.spi.function.SqlType;
+import com.facebook.presto.spi.function.TypeParameter;
+
+public class NativeMapSqlFunctions
+{
+    private NativeMapSqlFunctions() {}
+
+    @SqlInvokedScalarFunction(value = "map_top_n_keys", deterministic = true, calledOnNullInput = true)
+    @Description("Returns the top N keys of the given map sorting its keys using the provided lambda comparator.")
+    @TypeParameter("K")
+    @TypeParameter("V")
+    @SqlParameters({@SqlParameter(name = "input", type = "map(K, V)"), @SqlParameter(name = "n", type = "bigint"), @SqlParameter(name = "f", type = "function(K, K, bigint)")})
+    @SqlType("array")
+    public static String mapTopNKeysComparator()
+    {
+        return "RETURN IF(n < 0, fail('n must be greater than or equal to 0'), slice(reverse(array_sort(map_keys(input), f)), 1, n))";
+    }
+
+    @SqlInvokedScalarFunction(value = "map_top_n_values", deterministic = true, calledOnNullInput = true)
+    @Description("Returns the top N values of the given map sorted using the provided lambda comparator.")
+    @TypeParameter("K")
+    @TypeParameter("V")
+    @SqlParameters({@SqlParameter(name = "input", type = "map(K, V)"), @SqlParameter(name = "n", type = "bigint"), @SqlParameter(name = "f", type = "function(V, V, bigint)")})
+    @SqlType("array")
+    public static String mapTopNValuesComparator()
+    {
+        return "RETURN IF(n < 0, fail('n must be greater than or equal to 0'), slice(reverse(array_sort(remove_nulls(map_values(input)), f)) || filter(map_values(input), x -> x is null), 1, n))";
+    }
+}
diff --git a/presto-native-sql-invoked-functions-plugin/src/main/java/com/facebook/presto/scalar/sql/NativeSimpleSamplingPercent.java b/presto-native-sql-invoked-functions-plugin/src/main/java/com/facebook/presto/scalar/sql/NativeSimpleSamplingPercent.java
new file mode 100644
index 0000000000000..a710391760714
--- /dev/null
+++ b/presto-native-sql-invoked-functions-plugin/src/main/java/com/facebook/presto/scalar/sql/NativeSimpleSamplingPercent.java
@@ -0,0 +1,33 @@
+/*
+ * Licensed 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 com.facebook.presto.scalar.sql;
+
+import com.facebook.presto.spi.function.Description;
+import com.facebook.presto.spi.function.SqlInvokedScalarFunction;
+import com.facebook.presto.spi.function.SqlParameter;
+import com.facebook.presto.spi.function.SqlType;
+
+public class NativeSimpleSamplingPercent
+{
+    private NativeSimpleSamplingPercent() {}
+
+    @SqlInvokedScalarFunction(value = "key_sampling_percent", deterministic = true, calledOnNullInput = false)
+    @Description("Returns a value between 0.0 and 1.0 using the hash of the given input string")
+    @SqlParameter(name = "input", type = "varchar")
+    @SqlType("double")
+    public static String keySamplingPercent()
+    {
+        return "return (abs(from_ieee754_64(xxhash64(cast(input as varbinary)))) % 100) / 100. ";
+    }
+}
diff --git a/presto-native-sql-invoked-functions-plugin/src/main/java/com/facebook/presto/scalar/sql/NativeSqlInvokedFunctionsPlugin.java b/presto-native-sql-invoked-functions-plugin/src/main/java/com/facebook/presto/scalar/sql/NativeSqlInvokedFunctionsPlugin.java
new file mode 100644
index 0000000000000..69d7ff1e78522
--- /dev/null
+++ b/presto-native-sql-invoked-functions-plugin/src/main/java/com/facebook/presto/scalar/sql/NativeSqlInvokedFunctionsPlugin.java
@@ -0,0 +1,33 @@
+/*
+ * Licensed 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 com.facebook.presto.scalar.sql;
+
+import com.facebook.presto.spi.Plugin;
+import com.google.common.collect.ImmutableSet;
+
+import java.util.Set;
+
+public class NativeSqlInvokedFunctionsPlugin
+        implements Plugin
+{
+    @Override
+    public Set> getSqlInvokedFunctions()
+    {
+        return ImmutableSet.>builder()
+                .add(NativeArraySqlFunctions.class)
+                .add(NativeMapSqlFunctions.class)
+                .add(NativeSimpleSamplingPercent.class)
+                .build();
+    }
+}
diff --git a/presto-native-tests/pom.xml b/presto-native-tests/pom.xml
index e90785cbb6940..7e23670117b1d 100644
--- a/presto-native-tests/pom.xml
+++ b/presto-native-tests/pom.xml
@@ -192,6 +192,13 @@
             units
             test
         
+
+        
+            com.facebook.presto
+            presto-native-sql-invoked-functions-plugin
+            ${project.version}
+            test
+        
     
 
     
diff --git a/presto-native-tests/src/test/java/com/facebook/presto/nativetests/AbstractTestAggregationsNative.java b/presto-native-tests/src/test/java/com/facebook/presto/nativetests/AbstractTestAggregationsNative.java
index 39f787a1eb601..d668c65f098ea 100644
--- a/presto-native-tests/src/test/java/com/facebook/presto/nativetests/AbstractTestAggregationsNative.java
+++ b/presto-native-tests/src/test/java/com/facebook/presto/nativetests/AbstractTestAggregationsNative.java
@@ -24,13 +24,15 @@ public abstract class AbstractTestAggregationsNative
     private static final String QDIGEST_TYPE = "qdigest";
 
     private String storageFormat;
+    private boolean implicitCastCharNToVarchar;
     private String approxDistinctUnsupportedSignatureError;
     private String charTypeUnsupportedError;
     private String timeTypeUnsupportedError;
 
-    public void init(String storageFormat, boolean sidecarEnabled)
+    public void init(String storageFormat, boolean charNToVarcharImplicitCast, boolean sidecarEnabled)
     {
         this.storageFormat = storageFormat;
+        this.implicitCastCharNToVarchar = charNToVarcharImplicitCast;
         if (sidecarEnabled) {
             charTypeUnsupportedError = ".*Unknown type: char.*";
             timeTypeUnsupportedError = ".*Unknown type: time.*";
@@ -117,8 +119,16 @@ public void testApproximateCountDistinct()
         assertQuery("SELECT approx_distinct(CAST(custkey AS VARCHAR), 0.023) FROM orders", "SELECT 1036");
 
         // test char
-        assertQueryFails("SELECT approx_distinct(CAST(CAST(custkey AS VARCHAR) AS CHAR(20))) FROM orders", charTypeUnsupportedError, true);
-        assertQueryFails("SELECT approx_distinct(CAST(CAST(custkey AS VARCHAR) AS CHAR(20)), 0.023) FROM orders", charTypeUnsupportedError, true);
+        String charQuery = "SELECT approx_distinct(CAST(CAST(custkey AS VARCHAR) AS CHAR(20))) FROM orders";
+        String charWithErrorQuery = "SELECT approx_distinct(CAST(CAST(custkey AS VARCHAR) AS CHAR(20)), 0.023) FROM orders";
+        if (implicitCastCharNToVarchar) {
+            assertQuery(charQuery, "SELECT 1036");
+            assertQuery(charWithErrorQuery, "SELECT 1036");
+        }
+        else {
+            assertQueryFails(charQuery, charTypeUnsupportedError, true);
+            assertQueryFails(charWithErrorQuery, charTypeUnsupportedError, true);
+        }
 
         // test varbinary
         assertQuery("SELECT approx_distinct(to_utf8(CAST(custkey AS VARCHAR))) FROM orders", "SELECT 1036");
@@ -136,10 +146,13 @@ public void testSumDataSizeForStats()
         assertQuery("SELECT \"sum_data_size_for_stats\"(comment) FROM orders", "SELECT 787364");
 
         // char
-        // Presto removes trailing whitespaces when casting to CHAR.
-        // Hard code the expected data size since there is no easy to way to compute it in H2.
-        assertQueryFails("SELECT \"sum_data_size_for_stats\"(CAST(comment AS CHAR(1000))) FROM orders",
-                charTypeUnsupportedError, true);
+        String charQuery = "SELECT \"sum_data_size_for_stats\"(CAST(comment AS CHAR(1000))) FROM orders";
+        if (implicitCastCharNToVarchar) {
+            assertQuery(charQuery, "SELECT 787364");
+        }
+        else {
+            assertQueryFails(charQuery, charTypeUnsupportedError, true);
+        }
 
         // varbinary
         assertQuery("SELECT \"sum_data_size_for_stats\"(CAST(comment AS VARBINARY)) FROM orders", "SELECT 787364");
@@ -168,8 +181,13 @@ public void testMaxDataSizeForStats()
         assertQuery("SELECT \"max_data_size_for_stats\"(comment) FROM orders", "select 82");
 
         // char
-        assertQueryFails("SELECT \"max_data_size_for_stats\"(CAST(comment AS CHAR(1000))) FROM orders",
-                charTypeUnsupportedError, true);
+        String charQuery = "SELECT \"max_data_size_for_stats\"(CAST(comment AS CHAR(1000))) FROM orders";
+        if (implicitCastCharNToVarchar) {
+            assertQuery(charQuery, "SELECT 82");
+        }
+        else {
+            assertQueryFails(charQuery, charTypeUnsupportedError, true);
+        }
 
         // varbinary
         assertQuery("SELECT \"max_data_size_for_stats\"(CAST(comment AS VARBINARY)) FROM orders", "select 82");
diff --git a/presto-native-tests/src/test/java/com/facebook/presto/nativetests/NativeTestsUtils.java b/presto-native-tests/src/test/java/com/facebook/presto/nativetests/NativeTestsUtils.java
index b4dbaa197b36e..12d4cabed0220 100644
--- a/presto-native-tests/src/test/java/com/facebook/presto/nativetests/NativeTestsUtils.java
+++ b/presto-native-tests/src/test/java/com/facebook/presto/nativetests/NativeTestsUtils.java
@@ -24,13 +24,14 @@ public class NativeTestsUtils
 {
     private NativeTestsUtils() {}
 
-    public static QueryRunner createNativeQueryRunner(String storageFormat, boolean sidecarEnabled)
+    public static QueryRunner createNativeQueryRunner(String storageFormat, boolean charNToVarcharImplicitCast, boolean sidecarEnabled)
             throws Exception
     {
         QueryRunner queryRunner = nativeHiveQueryRunnerBuilder()
                 .setStorageFormat(storageFormat)
                 .setAddStorageFormatToPath(true)
                 .setUseThrift(true)
+                .setImplicitCastCharNToVarchar(charNToVarcharImplicitCast)
                 .setCoordinatorSidecarEnabled(sidecarEnabled)
                 .build();
         if (sidecarEnabled) {
@@ -39,6 +40,12 @@ public static QueryRunner createNativeQueryRunner(String storageFormat, boolean
         return queryRunner;
     }
 
+    public static QueryRunner createNativeQueryRunner(String storageFormat, boolean sidecarEnabled)
+            throws Exception
+    {
+        return createNativeQueryRunner(storageFormat, false, sidecarEnabled);
+    }
+
     public static void createTables(String storageFormat)
     {
         try {
@@ -58,4 +65,11 @@ public static void createTables(String storageFormat)
             throw new RuntimeException(e);
         }
     }
+
+    // TODO: remove and directly return charNToVarcharImplicitCast after addressing Issue #25894 and adding support for Char(n) type
+    // to class NativeTypeManager for Sidecar.
+    public static boolean getCharNToVarcharImplicitCastForTest(boolean sidecarEnabled, boolean charNToVarcharImplicitCast)
+    {
+        return sidecarEnabled ? false : charNToVarcharImplicitCast;
+    }
 }
diff --git a/presto-native-tests/src/test/java/com/facebook/presto/nativetests/TestAggregations.java b/presto-native-tests/src/test/java/com/facebook/presto/nativetests/TestAggregations.java
index 89d31a9257b95..220d14db2336a 100644
--- a/presto-native-tests/src/test/java/com/facebook/presto/nativetests/TestAggregations.java
+++ b/presto-native-tests/src/test/java/com/facebook/presto/nativetests/TestAggregations.java
@@ -22,6 +22,7 @@ public class TestAggregations
         extends AbstractTestAggregationsNative
 {
     private String storageFormat;
+    private boolean charNToVarcharImplicitCast;
     private boolean sidecarEnabled;
 
     @BeforeClass
@@ -30,15 +31,17 @@ public void init()
             throws Exception
     {
         storageFormat = System.getProperty("storageFormat", "PARQUET");
-        sidecarEnabled = parseBoolean(System.getProperty("sidecarEnabled", "true"));
-        super.init(storageFormat, sidecarEnabled);
+        sidecarEnabled = parseBoolean(System.getProperty("sidecarEnabled", "false"));
+        charNToVarcharImplicitCast = NativeTestsUtils.getCharNToVarcharImplicitCastForTest(
+                sidecarEnabled, parseBoolean(System.getProperty("charNToVarcharImplicitCast", "false")));
+        super.init(storageFormat, charNToVarcharImplicitCast, sidecarEnabled);
         super.init();
     }
 
     @Override
     protected QueryRunner createQueryRunner() throws Exception
     {
-        return NativeTestsUtils.createNativeQueryRunner(storageFormat, sidecarEnabled);
+        return NativeTestsUtils.createNativeQueryRunner(storageFormat, charNToVarcharImplicitCast, sidecarEnabled);
     }
 
     @Override
diff --git a/presto-native-tests/src/test/java/com/facebook/presto/nativetests/TestDistributedEngineOnlyQueries.java b/presto-native-tests/src/test/java/com/facebook/presto/nativetests/TestDistributedEngineOnlyQueries.java
index 71c884e828554..b54a9711940b9 100644
--- a/presto-native-tests/src/test/java/com/facebook/presto/nativetests/TestDistributedEngineOnlyQueries.java
+++ b/presto-native-tests/src/test/java/com/facebook/presto/nativetests/TestDistributedEngineOnlyQueries.java
@@ -123,4 +123,16 @@ public void testKeyBasedSampling()
     public void testDefaultSamplingPercent()
     {
     }
+
+    @Override
+    @Test(enabled = false)
+    public void testLeftJoinWithArrayContainsCondition()
+    {
+    }
+
+    @Override
+    @Test(enabled = false)
+    public void testTry()
+    {
+    }
 }
diff --git a/presto-native-tests/src/test/java/com/facebook/presto/nativetests/TestOptimizeMixedDistinctAggregations.java b/presto-native-tests/src/test/java/com/facebook/presto/nativetests/TestOptimizeMixedDistinctAggregations.java
index 9ff17c46fc9ef..8c775f51dd69a 100644
--- a/presto-native-tests/src/test/java/com/facebook/presto/nativetests/TestOptimizeMixedDistinctAggregations.java
+++ b/presto-native-tests/src/test/java/com/facebook/presto/nativetests/TestOptimizeMixedDistinctAggregations.java
@@ -26,6 +26,7 @@ public class TestOptimizeMixedDistinctAggregations
         extends AbstractTestAggregationsNative
 {
     private String storageFormat;
+    private boolean charNToVarcharImplicitCast;
     private boolean sidecarEnabled;
 
     @BeforeClass
@@ -35,7 +36,9 @@ public void init()
     {
         storageFormat = System.getProperty("storageFormat", "PARQUET");
         sidecarEnabled = parseBoolean(System.getProperty("sidecarEnabled", "true"));
-        super.init(storageFormat, sidecarEnabled);
+        charNToVarcharImplicitCast = NativeTestsUtils.getCharNToVarcharImplicitCastForTest(
+                sidecarEnabled, parseBoolean(System.getProperty("charNToVarcharImplicitCast", "false")));
+        super.init(storageFormat, charNToVarcharImplicitCast, sidecarEnabled);
         super.init();
     }
 
@@ -45,6 +48,7 @@ protected QueryRunner createQueryRunner() throws Exception
         QueryRunner queryRunner = nativeHiveQueryRunnerBuilder()
                 .setStorageFormat(storageFormat)
                 .setAddStorageFormatToPath(true)
+                .setImplicitCastCharNToVarchar(charNToVarcharImplicitCast)
                 .setUseThrift(true)
                 .setCoordinatorSidecarEnabled(sidecarEnabled)
                 .setExtraCoordinatorProperties(
diff --git a/presto-native-tests/src/test/java/com/facebook/presto/nativetests/TestWindowQueries.java b/presto-native-tests/src/test/java/com/facebook/presto/nativetests/TestWindowQueries.java
index 6db7af6e559ed..3b2a2a8c44d05 100644
--- a/presto-native-tests/src/test/java/com/facebook/presto/nativetests/TestWindowQueries.java
+++ b/presto-native-tests/src/test/java/com/facebook/presto/nativetests/TestWindowQueries.java
@@ -13,18 +13,11 @@
  */
 package com.facebook.presto.nativetests;
 
-import com.facebook.presto.testing.MaterializedResult;
 import com.facebook.presto.testing.QueryRunner;
 import com.facebook.presto.tests.AbstractTestWindowQueries;
-import org.intellij.lang.annotations.Language;
 import org.testng.annotations.BeforeClass;
 import org.testng.annotations.Test;
 
-import static com.facebook.presto.common.type.BigintType.BIGINT;
-import static com.facebook.presto.common.type.VarcharType.createUnboundedVarcharType;
-import static com.facebook.presto.common.type.VarcharType.createVarcharType;
-import static com.facebook.presto.testing.MaterializedResult.resultBuilder;
-import static com.facebook.presto.testing.assertions.Assert.assertEquals;
 import static java.lang.Boolean.parseBoolean;
 
 public class TestWindowQueries
@@ -234,29 +227,4 @@ public void testTypes()
                         "(INTERVAL '2' month, ARRAY[INTERVAL '1' month, INTERVAL '2' month]), " +
                         "(INTERVAL '5' year, ARRAY[INTERVAL '5' year])");
     }
-
-    // Todo: Refactor this test case when support for varchar(N) is added in native execution.
-    // The return types do not match on the native query runner : types=[varchar, bigint] and the java query runner:  types=[varchar(3), bigint].
-    @Override
-    @Test
-    public void testWindowFunctionWithGroupBy()
-    {
-        @Language("SQL") String sql = "SELECT *, rank() OVER (PARTITION BY x)\n" +
-                "FROM (SELECT 'foo' x)\n" +
-                "GROUP BY 1";
-
-        MaterializedResult actual = computeActual(sql);
-        MaterializedResult expected;
-        if (sidecarEnabled) {
-            expected = resultBuilder(getSession(), createUnboundedVarcharType(), BIGINT)
-                    .row("foo", 1L)
-                    .build();
-        }
-        else {
-            expected = resultBuilder(getSession(), createVarcharType(3), BIGINT)
-                    .row("foo", 1L)
-                    .build();
-        }
-        assertEquals(actual, expected);
-    }
 }
diff --git a/presto-parser/src/main/java/com/facebook/presto/sql/ExpressionFormatter.java b/presto-parser/src/main/java/com/facebook/presto/sql/ExpressionFormatter.java
index bc03fb68154ea..bae8559c7f127 100644
--- a/presto-parser/src/main/java/com/facebook/presto/sql/ExpressionFormatter.java
+++ b/presto-parser/src/main/java/com/facebook/presto/sql/ExpressionFormatter.java
@@ -33,6 +33,7 @@
 import com.facebook.presto.sql.tree.DecimalLiteral;
 import com.facebook.presto.sql.tree.DereferenceExpression;
 import com.facebook.presto.sql.tree.DoubleLiteral;
+import com.facebook.presto.sql.tree.EnumLiteral;
 import com.facebook.presto.sql.tree.ExistsPredicate;
 import com.facebook.presto.sql.tree.Expression;
 import com.facebook.presto.sql.tree.Extract;
@@ -198,6 +199,12 @@ protected String visitBooleanLiteral(BooleanLiteral node, Void context)
             return String.valueOf(node.getValue());
         }
 
+        @Override
+        protected String visitEnumLiteral(EnumLiteral node, Void context)
+        {
+            return node.getType() + ": " + node.getValue();
+        }
+
         @Override
         protected String visitStringLiteral(StringLiteral node, Void context)
         {
diff --git a/presto-parser/src/main/java/com/facebook/presto/sql/tree/Expression.java b/presto-parser/src/main/java/com/facebook/presto/sql/tree/Expression.java
index 693852f8a8c39..48761444c2b29 100644
--- a/presto-parser/src/main/java/com/facebook/presto/sql/tree/Expression.java
+++ b/presto-parser/src/main/java/com/facebook/presto/sql/tree/Expression.java
@@ -39,4 +39,10 @@ public final String toString()
     {
         return ExpressionFormatter.formatExpression(this, Optional.empty()); // This will not replace parameters, but we don't have access to them here
     }
+
+    @Override
+    public String getArgumentTypeString()
+    {
+        return "expression";
+    }
 }
diff --git a/presto-parser/src/main/java/com/facebook/presto/sql/tree/Node.java b/presto-parser/src/main/java/com/facebook/presto/sql/tree/Node.java
index 2317c5258c1c9..b1c1354b09d8c 100644
--- a/presto-parser/src/main/java/com/facebook/presto/sql/tree/Node.java
+++ b/presto-parser/src/main/java/com/facebook/presto/sql/tree/Node.java
@@ -40,6 +40,11 @@ public Optional getLocation()
         return location;
     }
 
+    public String getArgumentTypeString()
+    {
+        throw new IllegalArgumentException("Unexpected table function argument type: " + this.getClass().getSimpleName());
+    }
+
     public abstract List getChildren();
 
     // Force subclasses to have a proper equals and hashcode implementation
diff --git a/presto-parser/src/main/java/com/facebook/presto/sql/tree/TableFunctionDescriptorArgument.java b/presto-parser/src/main/java/com/facebook/presto/sql/tree/TableFunctionDescriptorArgument.java
index 7c6150ee320d8..e762c99b0fa70 100644
--- a/presto-parser/src/main/java/com/facebook/presto/sql/tree/TableFunctionDescriptorArgument.java
+++ b/presto-parser/src/main/java/com/facebook/presto/sql/tree/TableFunctionDescriptorArgument.java
@@ -83,4 +83,10 @@ public String toString()
     {
         return descriptor.map(Descriptor::toString).orElse("CAST (NULL AS DESCRIPTOR)");
     }
+
+    @Override
+    public String getArgumentTypeString()
+    {
+        return "descriptor";
+    }
 }
diff --git a/presto-parser/src/main/java/com/facebook/presto/sql/tree/TableFunctionTableArgument.java b/presto-parser/src/main/java/com/facebook/presto/sql/tree/TableFunctionTableArgument.java
index b23da149cabae..ef0eb8bc20d96 100644
--- a/presto-parser/src/main/java/com/facebook/presto/sql/tree/TableFunctionTableArgument.java
+++ b/presto-parser/src/main/java/com/facebook/presto/sql/tree/TableFunctionTableArgument.java
@@ -120,4 +120,10 @@ public String toString()
 
         return builder.toString();
     }
+
+    @Override
+    public String getArgumentTypeString()
+    {
+        return "table";
+    }
 }
diff --git a/presto-pinot-toolkit/pom.xml b/presto-pinot-toolkit/pom.xml
index 88367dbd9fde3..5abe5d9005762 100644
--- a/presto-pinot-toolkit/pom.xml
+++ b/presto-pinot-toolkit/pom.xml
@@ -515,7 +515,7 @@
         
             org.reflections
             reflections
-            0.9.9
+            0.10.2
             
                 
                     com.google.code.findbugs
diff --git a/presto-pinot-toolkit/src/main/java/com/facebook/presto/pinot/PinotMetadata.java b/presto-pinot-toolkit/src/main/java/com/facebook/presto/pinot/PinotMetadata.java
index 6dd350c8e7491..09af2042494d4 100644
--- a/presto-pinot-toolkit/src/main/java/com/facebook/presto/pinot/PinotMetadata.java
+++ b/presto-pinot-toolkit/src/main/java/com/facebook/presto/pinot/PinotMetadata.java
@@ -80,7 +80,7 @@ public PinotTableHandle getTableHandle(ConnectorSession session, SchemaTableName
     }
 
     @Override
-    public List getTableLayouts(
+    public ConnectorTableLayoutResult getTableLayoutForConstraint(
             ConnectorSession session,
             ConnectorTableHandle table,
             Constraint constraint,
@@ -89,7 +89,7 @@ public List getTableLayouts(
         // Constraint's don't need to be pushed down since they are already taken care off by the pushdown logic
         PinotTableHandle pinotTableHandle = (PinotTableHandle) table;
         ConnectorTableLayout layout = new ConnectorTableLayout(new PinotTableLayoutHandle(pinotTableHandle));
-        return ImmutableList.of(new ConnectorTableLayoutResult(layout, constraint.getSummary()));
+        return new ConnectorTableLayoutResult(layout, constraint.getSummary());
     }
 
     @Override
diff --git a/presto-pinot/src/test/java/com/facebook/presto/pinot/udf/TestPinotFunctions.java b/presto-pinot/src/test/java/com/facebook/presto/pinot/udf/TestPinotFunctions.java
index acfd681d0e48b..3dec94b4fe4ce 100644
--- a/presto-pinot/src/test/java/com/facebook/presto/pinot/udf/TestPinotFunctions.java
+++ b/presto-pinot/src/test/java/com/facebook/presto/pinot/udf/TestPinotFunctions.java
@@ -15,15 +15,23 @@
 
 import com.facebook.presto.operator.scalar.AbstractTestFunctions;
 import com.facebook.presto.pinot.PinotPlugin;
+import com.facebook.presto.sql.analyzer.FeaturesConfig;
+import com.facebook.presto.sql.analyzer.FunctionsConfig;
 import org.testng.annotations.BeforeClass;
 import org.testng.annotations.Test;
 
+import static com.facebook.presto.SessionTestUtils.TEST_SESSION;
 import static com.facebook.presto.common.type.DoubleType.DOUBLE;
 import static com.facebook.presto.metadata.FunctionExtractor.extractFunctions;
 
 public class TestPinotFunctions
         extends AbstractTestFunctions
 {
+    public TestPinotFunctions()
+    {
+        super(TEST_SESSION, new FeaturesConfig(), new FunctionsConfig(), false);
+    }
+
     @BeforeClass
     public void setUp()
     {
diff --git a/presto-plan-checker-router-plugin/pom.xml b/presto-plan-checker-router-plugin/pom.xml
index 60e34e4147d19..db1c4b0f1736c 100644
--- a/presto-plan-checker-router-plugin/pom.xml
+++ b/presto-plan-checker-router-plugin/pom.xml
@@ -223,6 +223,13 @@
             presto-hive-metastore
             test
         
+
+        
+            com.facebook.presto
+            presto-native-sql-invoked-functions-plugin
+            ${project.version}
+            test
+        
     
 
     
diff --git a/presto-product-tests/src/main/java/com/facebook/presto/tests/SystemConnectorTests.java b/presto-product-tests/src/main/java/com/facebook/presto/tests/SystemConnectorTests.java
index 3317d07e90f8c..4f14b99e2d931 100644
--- a/presto-product-tests/src/main/java/com/facebook/presto/tests/SystemConnectorTests.java
+++ b/presto-product-tests/src/main/java/com/facebook/presto/tests/SystemConnectorTests.java
@@ -108,9 +108,9 @@ public void selectRuntimeTasks()
     @Test(groups = {SYSTEM_CONNECTOR, JDBC})
     public void selectMetadataCatalogs()
     {
-        String sql = "select catalog_name, connector_id from system.metadata.catalogs";
+        String sql = "select catalog_name, connector_id, connector_name from system.metadata.catalogs";
         assertThat(query(sql))
-                .hasColumns(VARCHAR, VARCHAR)
+                .hasColumns(VARCHAR, VARCHAR, VARCHAR)
                 .hasAnyRows();
     }
 
diff --git a/presto-product-tests/src/main/resources/sql-tests/testcases/system/selectInformationSchemaColumns.result b/presto-product-tests/src/main/resources/sql-tests/testcases/system/selectInformationSchemaColumns.result
index c442a420ebead..4103c03301025 100644
--- a/presto-product-tests/src/main/resources/sql-tests/testcases/system/selectInformationSchemaColumns.result
+++ b/presto-product-tests/src/main/resources/sql-tests/testcases/system/selectInformationSchemaColumns.result
@@ -46,6 +46,7 @@ system| metadata| analyze_properties| type| varchar| YES| null| null|
 system| metadata| analyze_properties| description| varchar| YES| null| null|
 system| metadata| catalogs| catalog_name| varchar| YES| null| null|
 system| metadata| catalogs| connector_id| varchar| YES| null| null|
+system| metadata| catalogs| connector_name| varchar| YES| null| null|
 system| metadata| column_properties| catalog_name| varchar| YES| null| null|
 system| metadata| column_properties| property_name| varchar| YES| null| null|
 system| metadata| column_properties| default_value| varchar| YES| null| null|
diff --git a/presto-prometheus/src/main/java/com/facebook/presto/plugin/prometheus/PrometheusMetadata.java b/presto-prometheus/src/main/java/com/facebook/presto/plugin/prometheus/PrometheusMetadata.java
index bb223c7836421..07b68eee5cb1f 100644
--- a/presto-prometheus/src/main/java/com/facebook/presto/plugin/prometheus/PrometheusMetadata.java
+++ b/presto-prometheus/src/main/java/com/facebook/presto/plugin/prometheus/PrometheusMetadata.java
@@ -77,11 +77,15 @@ public ConnectorTableHandle getTableHandle(ConnectorSession session, SchemaTable
     }
 
     @Override
-    public List getTableLayouts(ConnectorSession session, ConnectorTableHandle table, Constraint constraint, Optional> desiredColumns)
+    public ConnectorTableLayoutResult getTableLayoutForConstraint(
+            ConnectorSession session,
+            ConnectorTableHandle table,
+            Constraint constraint,
+            Optional> desiredColumns)
     {
         PrometheusTableHandle tableHandle = (PrometheusTableHandle) table;
         ConnectorTableLayout layout = new ConnectorTableLayout(new PrometheusTableLayoutHandle(tableHandle, constraint.getSummary()));
-        return ImmutableList.of(new ConnectorTableLayoutResult(layout, constraint.getSummary()));
+        return new ConnectorTableLayoutResult(layout, constraint.getSummary());
     }
 
     @Override
diff --git a/presto-redis/src/main/java/com/facebook/presto/redis/RedisMetadata.java b/presto-redis/src/main/java/com/facebook/presto/redis/RedisMetadata.java
index acb4fef35a7e8..8295dd7d0d90c 100644
--- a/presto-redis/src/main/java/com/facebook/presto/redis/RedisMetadata.java
+++ b/presto-redis/src/main/java/com/facebook/presto/redis/RedisMetadata.java
@@ -124,7 +124,7 @@ public ConnectorTableMetadata getTableMetadata(ConnectorSession session, Connect
     }
 
     @Override
-    public List getTableLayouts(
+    public ConnectorTableLayoutResult getTableLayoutForConstraint(
             ConnectorSession session,
             ConnectorTableHandle table,
             Constraint constraint,
@@ -134,7 +134,7 @@ public List getTableLayouts(
 
         ConnectorTableLayout layout = new ConnectorTableLayout(new RedisTableLayoutHandle(tableHandle));
 
-        return ImmutableList.of(new ConnectorTableLayoutResult(layout, constraint.getSummary()));
+        return new ConnectorTableLayoutResult(layout, constraint.getSummary());
     }
 
     @Override
@@ -143,8 +143,7 @@ public ConnectorTableLayout getTableLayout(ConnectorSession session, ConnectorTa
         RedisTableLayoutHandle layout = convertLayout(handle);
 
         // tables in this connector have a single layout
-        return getTableLayouts(session, layout.getTable(), Constraint.alwaysTrue(), Optional.empty())
-                .get(0)
+        return getTableLayoutForConstraint(session, layout.getTable(), Constraint.alwaysTrue(), Optional.empty())
                 .getTableLayout();
     }
 
diff --git a/presto-redis/src/test/java/com/facebook/presto/redis/util/RedisLoader.java b/presto-redis/src/test/java/com/facebook/presto/redis/util/RedisLoader.java
index b33f6e14fcf7f..52b51368db4e3 100644
--- a/presto-redis/src/test/java/com/facebook/presto/redis/util/RedisLoader.java
+++ b/presto-redis/src/test/java/com/facebook/presto/redis/util/RedisLoader.java
@@ -143,7 +143,7 @@ public void addResults(QueryStatusInfo statusInfo, QueryData data)
         }
 
         @Override
-        public Void build(Map setSessionProperties, Set resetSessionProperties)
+        public Void build(Map setSessionProperties, Set resetSessionProperties, String startTransactionId, boolean clearTransactionId)
         {
             return null;
         }
diff --git a/presto-resource-group-managers/src/test/java/com/facebook/presto/resourceGroups/reloading/TestReloadingResourceGroupConfigurationManager.java b/presto-resource-group-managers/src/test/java/com/facebook/presto/resourceGroups/reloading/TestReloadingResourceGroupConfigurationManager.java
index 8f02abab278e5..065fbfd740daf 100644
--- a/presto-resource-group-managers/src/test/java/com/facebook/presto/resourceGroups/reloading/TestReloadingResourceGroupConfigurationManager.java
+++ b/presto-resource-group-managers/src/test/java/com/facebook/presto/resourceGroups/reloading/TestReloadingResourceGroupConfigurationManager.java
@@ -15,8 +15,12 @@
 
 import com.facebook.airlift.units.DataSize;
 import com.facebook.airlift.units.Duration;
+import com.facebook.presto.execution.ClusterOverloadConfig;
 import com.facebook.presto.execution.resourceGroups.InternalResourceGroup;
+import com.facebook.presto.execution.scheduler.clusterOverload.ClusterOverloadPolicy;
+import com.facebook.presto.execution.scheduler.clusterOverload.ClusterResourceChecker;
 import com.facebook.presto.metadata.InMemoryNodeManager;
+import com.facebook.presto.metadata.InternalNodeManager;
 import com.facebook.presto.resourceGroups.VariableMap;
 import com.facebook.presto.resourceGroups.db.DbManagerSpecProvider;
 import com.facebook.presto.resourceGroups.db.DbResourceGroupConfig;
@@ -74,7 +78,7 @@ public void testConfiguration()
         DbManagerSpecProvider dbManagerSpecProvider = new DbManagerSpecProvider(daoProvider.get(), ENVIRONMENT, new ReloadingResourceGroupConfig());
         ReloadingResourceGroupConfigurationManager manager = new ReloadingResourceGroupConfigurationManager((poolId, listener) -> {}, new ReloadingResourceGroupConfig(), dbManagerSpecProvider);
         AtomicBoolean exported = new AtomicBoolean();
-        InternalResourceGroup global = new InternalResourceGroup.RootInternalResourceGroup("global", (group, export) -> exported.set(export), directExecutor(), ignored -> Optional.empty(), rg -> false, new InMemoryNodeManager());
+        InternalResourceGroup global = new InternalResourceGroup.RootInternalResourceGroup("global", (group, export) -> exported.set(export), directExecutor(), ignored -> Optional.empty(), rg -> false, new InMemoryNodeManager(), createClusterResourceChecker());
         manager.configure(global, new SelectionContext<>(global.getId(), new VariableMap(ImmutableMap.of("USER", "user"))));
         assertEqualsResourceGroup(global, "1MB", 1000, 100, 100, WEIGHTED, DEFAULT_WEIGHT, true, new Duration(1, HOURS), new Duration(1, DAYS), new ResourceGroupQueryLimits(Optional.of(new Duration(1, HOURS)), Optional.of(new DataSize(1, MEGABYTE)), Optional.of(new Duration(1, HOURS))));
         exported.set(false);
@@ -97,7 +101,7 @@ public void testMissing()
         dao.insertSelector(2, 1, null, null, null, null, null, null);
         DbManagerSpecProvider dbManagerSpecProvider = new DbManagerSpecProvider(daoProvider.get(), ENVIRONMENT, new ReloadingResourceGroupConfig());
         ReloadingResourceGroupConfigurationManager manager = new ReloadingResourceGroupConfigurationManager((poolId, listener) -> {}, new ReloadingResourceGroupConfig(), dbManagerSpecProvider);
-        InternalResourceGroup missing = new InternalResourceGroup.RootInternalResourceGroup("missing", (group, export) -> {}, directExecutor(), ignored -> Optional.empty(), rg -> false, new InMemoryNodeManager());
+        InternalResourceGroup missing = new InternalResourceGroup.RootInternalResourceGroup("missing", (group, export) -> {}, directExecutor(), ignored -> Optional.empty(), rg -> false, new InMemoryNodeManager(), createClusterResourceChecker());
         manager.configure(missing, new SelectionContext<>(missing.getId(), new VariableMap(ImmutableMap.of("USER", "user"))));
     }
 
@@ -118,7 +122,7 @@ public void testReconfig()
         ReloadingResourceGroupConfigurationManager manager = new ReloadingResourceGroupConfigurationManager((poolId, listener) -> {}, new ReloadingResourceGroupConfig(), dbManagerSpecProvider);
         manager.start();
         AtomicBoolean exported = new AtomicBoolean();
-        InternalResourceGroup global = new InternalResourceGroup.RootInternalResourceGroup("global", (group, export) -> exported.set(export), directExecutor(), ignored -> Optional.empty(), rg -> false, new InMemoryNodeManager());
+        InternalResourceGroup global = new InternalResourceGroup.RootInternalResourceGroup("global", (group, export) -> exported.set(export), directExecutor(), ignored -> Optional.empty(), rg -> false, new InMemoryNodeManager(), createClusterResourceChecker());
         manager.configure(global, new SelectionContext<>(global.getId(), new VariableMap(ImmutableMap.of("USER", "user"))));
         InternalResourceGroup globalSub = global.getOrCreateSubGroup("sub", true);
         manager.configure(globalSub, new SelectionContext<>(globalSub.getId(), new VariableMap(ImmutableMap.of("USER", "user"))));
@@ -247,4 +251,29 @@ private static void assertEqualsResourceGroup(
         assertEquals(group.getHardCpuLimit(), hardCpuLimit);
         assertEquals(group.getPerQueryLimits(), perQueryLimits);
     }
+
+    private static ClusterResourceChecker createClusterResourceChecker()
+    {
+        // Create a mock cluster overload policy that never reports overload
+        ClusterOverloadPolicy mockPolicy = new ClusterOverloadPolicy()
+        {
+            @Override
+            public boolean isClusterOverloaded(InternalNodeManager nodeManager)
+            {
+                return false; // Never overloaded for tests
+            }
+
+            @Override
+            public String getName()
+            {
+                return "test-policy";
+            }
+        };
+
+        // Create a config with throttling disabled for tests
+        ClusterOverloadConfig config = new ClusterOverloadConfig()
+                .setClusterOverloadThrottlingEnabled(false);
+
+        return new ClusterResourceChecker(mockPolicy, config, new InMemoryNodeManager());
+    }
 }
diff --git a/presto-router/src/main/java/com/facebook/presto/router/RouterModule.java b/presto-router/src/main/java/com/facebook/presto/router/RouterModule.java
index efc7cc783ddb7..168668fe6ed3f 100755
--- a/presto-router/src/main/java/com/facebook/presto/router/RouterModule.java
+++ b/presto-router/src/main/java/com/facebook/presto/router/RouterModule.java
@@ -15,6 +15,7 @@
 
 import com.facebook.airlift.configuration.AbstractConfigurationAwareModule;
 import com.facebook.airlift.units.Duration;
+import com.facebook.presto.ClientRequestFilterManager;
 import com.facebook.presto.client.NodeVersion;
 import com.facebook.presto.router.cluster.ClusterManager;
 import com.facebook.presto.router.cluster.ClusterManager.ClusterStatusTracker;
@@ -70,6 +71,7 @@ protected void setup(Binder binder)
     {
         ServerConfig serverConfig = buildConfigObject(ServerConfig.class);
 
+        binder.bind(ClientRequestFilterManager.class).in(Scopes.SINGLETON);
         binder.bind(RouterPluginManager.class).in(Scopes.SINGLETON);
         configBinder(binder).bindConfig(RouterConfig.class);
 
diff --git a/presto-server/src/main/provisio/presto.xml b/presto-server/src/main/provisio/presto.xml
index 14bcfd5b3c31a..b7176785fedd4 100644
--- a/presto-server/src/main/provisio/presto.xml
+++ b/presto-server/src/main/provisio/presto.xml
@@ -28,7 +28,7 @@
         
     
 
-    
+    
     
         
             
@@ -41,8 +41,14 @@
         
     
 
-    
-        
+    
+        
+            
+        
+    
+
+    
+        
             
         
     
@@ -275,9 +281,100 @@
         
     
 
-    
+    
+        
+            
+        
+    
+
+    
+    
+        
+            
+        
+    
+
+    
+        
+            
+        
+    
+
+    
+        
+            
+        
+    
+
+    
+        
+            
+        
+    
+
+    
+        
+            
+        
+    
+
+    
+        
+            
+        
+    
+
+    
+        
+            
+        
+    
+
+    
+        
+            
+        
+    
+
+    
+        
+            
+        
+    
+
+    
+        
+            
+        
+    
+
+    
+        
+            
+        
+    
+
+    
+        
+            
+        
+    
+
+    
+        
+            
+        
+    
+
+    
         
             
         
     
+
+    
+        
+            
+        
+    
 
diff --git a/presto-session-property-managers-common/pom.xml b/presto-session-property-managers-common/pom.xml
new file mode 100644
index 0000000000000..098b68e7ea58c
--- /dev/null
+++ b/presto-session-property-managers-common/pom.xml
@@ -0,0 +1,81 @@
+
+
+    4.0.0
+
+    
+        com.facebook.presto
+        presto-root
+        0.295-SNAPSHOT
+    
+
+    presto-session-property-managers-common
+    presto-session-property-managers-common
+    Presto - Session Property Managers Common
+    jar
+
+    
+        ${project.parent.basedir}
+    
+
+    
+        
+            com.google.guava
+            guava
+        
+
+        
+            org.jdbi
+            jdbi3-core
+        
+
+        
+        
+            com.facebook.presto
+            presto-spi
+            provided
+        
+
+        
+            com.facebook.presto
+            presto-common
+            provided
+        
+
+        
+            com.facebook.airlift
+            units
+            provided
+        
+
+        
+            io.airlift
+            slice
+            provided
+        
+
+        
+            com.fasterxml.jackson.core
+            jackson-annotations
+            provided
+        
+
+        
+            org.openjdk.jol
+            jol-core
+            provided
+        
+
+        
+        
+            com.facebook.presto
+            presto-testng-services
+            test
+        
+
+        
+            org.testng
+            testng
+            test
+        
+    
+
diff --git a/presto-session-property-managers-common/src/main/java/com/facebook/presto/session/AbstractSessionPropertyManager.java b/presto-session-property-managers-common/src/main/java/com/facebook/presto/session/AbstractSessionPropertyManager.java
new file mode 100644
index 0000000000000..7780652db038a
--- /dev/null
+++ b/presto-session-property-managers-common/src/main/java/com/facebook/presto/session/AbstractSessionPropertyManager.java
@@ -0,0 +1,65 @@
+/*
+ * Licensed 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 com.facebook.presto.session;
+
+import com.facebook.presto.spi.session.SessionConfigurationContext;
+import com.facebook.presto.spi.session.SessionPropertyConfigurationManager;
+
+import java.util.HashMap;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.Set;
+
+public abstract class AbstractSessionPropertyManager
+        implements SessionPropertyConfigurationManager
+{
+    private static final String DEFAULT_PROPERTIES = "defaultProperties";
+    private static final String OVERRIDE_PROPERTIES = "overrideProperties";
+
+    @Override
+    public final SystemSessionPropertyConfiguration getSystemSessionProperties(SessionConfigurationContext context)
+    {
+        Map defaultProperties = new HashMap<>();
+        Set overridePropertyNames = new HashSet<>();
+        for (SessionMatchSpec sessionMatchSpec : getSessionMatchSpecs()) {
+            Map newProperties = sessionMatchSpec.match(sessionMatchSpec.getSessionProperties(), context);
+            defaultProperties.putAll(newProperties);
+            if (sessionMatchSpec.getOverrideSessionProperties().orElse(false)) {
+                overridePropertyNames.addAll(newProperties.keySet());
+            }
+        }
+
+        // Once a property has been overridden it stays that way and the value is updated by any rule
+        Map overrideProperties = new HashMap<>();
+        for (String propertyName : overridePropertyNames) {
+            overrideProperties.put(propertyName, defaultProperties.get(propertyName));
+        }
+
+        return new SystemSessionPropertyConfiguration(defaultProperties, overrideProperties);
+    }
+
+    @Override
+    public final Map> getCatalogSessionProperties(SessionConfigurationContext context)
+    {
+        Map> catalogProperties = new HashMap<>();
+        for (SessionMatchSpec sessionMatchSpec : getSessionMatchSpecs()) {
+            Map> newProperties = sessionMatchSpec.match(sessionMatchSpec.getCatalogSessionProperties(), context);
+            catalogProperties.putAll(newProperties);
+        }
+        return catalogProperties;
+    }
+
+    protected abstract List getSessionMatchSpecs();
+}
diff --git a/presto-session-property-managers-common/src/main/java/com/facebook/presto/session/SessionMatchSpec.java b/presto-session-property-managers-common/src/main/java/com/facebook/presto/session/SessionMatchSpec.java
new file mode 100644
index 0000000000000..6a70aaae42d69
--- /dev/null
+++ b/presto-session-property-managers-common/src/main/java/com/facebook/presto/session/SessionMatchSpec.java
@@ -0,0 +1,252 @@
+/*
+ * Licensed 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 com.facebook.presto.session;
+
+import com.facebook.presto.spi.resourceGroups.ResourceGroupId;
+import com.facebook.presto.spi.session.SessionConfigurationContext;
+import com.fasterxml.jackson.annotation.JsonCreator;
+import com.fasterxml.jackson.annotation.JsonProperty;
+import com.google.common.base.Splitter;
+import com.google.common.collect.ImmutableList;
+import com.google.common.collect.ImmutableMap;
+import com.google.common.collect.ImmutableSet;
+import org.jdbi.v3.core.mapper.RowMapper;
+import org.jdbi.v3.core.statement.StatementContext;
+
+import java.sql.ResultSet;
+import java.sql.SQLException;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+import java.util.Optional;
+import java.util.Set;
+import java.util.regex.Pattern;
+
+import static com.google.common.base.Preconditions.checkArgument;
+import static java.util.Objects.requireNonNull;
+
+public class SessionMatchSpec
+{
+    public static String emptyCatalog = "__NULL__";
+    private final Optional userRegex;
+    private final Optional sourceRegex;
+    private final Set clientTags;
+    private final Optional queryType;
+    private final Optional clientInfoRegex;
+    private final Optional resourceGroupRegex;
+    private final Optional overrideSessionProperties;
+    private final Map sessionProperties;
+    private final Map> catalogSessionProperties;
+
+    @JsonCreator
+    public SessionMatchSpec(
+            @JsonProperty("user") Optional userRegex,
+            @JsonProperty("source") Optional sourceRegex,
+            @JsonProperty("clientTags") Optional> clientTags,
+            @JsonProperty("queryType") Optional queryType,
+            @JsonProperty("group") Optional resourceGroupRegex,
+            @JsonProperty("clientInfo") Optional clientInfoRegex,
+            @JsonProperty("overrideSessionProperties") Optional overrideSessionProperties,
+            @JsonProperty("sessionProperties") Map sessionProperties,
+            @JsonProperty("catalogSessionProperties") Map> catalogSessionProperties)
+    {
+        this.userRegex = requireNonNull(userRegex, "userRegex is null");
+        this.sourceRegex = requireNonNull(sourceRegex, "sourceRegex is null");
+        requireNonNull(clientTags, "clientTags is null");
+        this.clientTags = ImmutableSet.copyOf(clientTags.orElse(ImmutableList.of()));
+        this.queryType = requireNonNull(queryType, "queryType is null");
+        this.resourceGroupRegex = requireNonNull(resourceGroupRegex, "resourceGroupRegex is null");
+        this.clientInfoRegex = requireNonNull(clientInfoRegex, "clientInfoRegex is null");
+        this.overrideSessionProperties = requireNonNull(overrideSessionProperties, "overrideSessionProperties is null");
+
+        checkArgument(sessionProperties != null || catalogSessionProperties != null,
+                "Either sessionProperties or catalogSessionProperties must be provided");
+
+        this.sessionProperties = ImmutableMap.copyOf(Optional.ofNullable(sessionProperties).orElse(ImmutableMap.of()));
+        this.catalogSessionProperties = ImmutableMap.copyOf(Optional.ofNullable(catalogSessionProperties).orElse(ImmutableMap.of()));
+    }
+
+    public  Map match(Map object, SessionConfigurationContext context)
+    {
+        if (userRegex.isPresent() && !userRegex.get().matcher(context.getUser()).matches()) {
+            return ImmutableMap.of();
+        }
+        if (sourceRegex.isPresent()) {
+            String source = context.getSource().orElse("");
+            if (!sourceRegex.get().matcher(source).matches()) {
+                return ImmutableMap.of();
+            }
+        }
+        if (!clientTags.isEmpty() && !context.getClientTags().containsAll(clientTags)) {
+            return ImmutableMap.of();
+        }
+
+        if (queryType.isPresent()) {
+            String contextQueryType = context.getQueryType().orElse("");
+            if (!queryType.get().equalsIgnoreCase(contextQueryType)) {
+                return ImmutableMap.of();
+            }
+        }
+
+        if (clientInfoRegex.isPresent()) {
+            String clientInfo = context.getClientInfo().orElse("");
+            if (!clientInfoRegex.get().matcher(clientInfo).matches()) {
+                return ImmutableMap.of();
+            }
+        }
+
+        if (resourceGroupRegex.isPresent()) {
+            String resourceGroupId = context.getResourceGroupId().map(ResourceGroupId::toString).orElse("");
+            if (!resourceGroupRegex.get().matcher(resourceGroupId).matches()) {
+                return ImmutableMap.of();
+            }
+        }
+
+        return object;
+    }
+
+    @JsonProperty("user")
+    public Optional getUserRegex()
+    {
+        return userRegex;
+    }
+
+    @JsonProperty("source")
+    public Optional getSourceRegex()
+    {
+        return sourceRegex;
+    }
+
+    @JsonProperty
+    public Set getClientTags()
+    {
+        return clientTags;
+    }
+
+    @JsonProperty
+    public Optional getQueryType()
+    {
+        return queryType;
+    }
+
+    @JsonProperty("group")
+    public Optional getResourceGroupRegex()
+    {
+        return resourceGroupRegex;
+    }
+
+    @JsonProperty
+    public Optional getClientInfoRegex()
+    {
+        return clientInfoRegex;
+    }
+
+    @JsonProperty
+    public Optional getOverrideSessionProperties()
+    {
+        return overrideSessionProperties;
+    }
+
+    @JsonProperty
+    public Map getSessionProperties()
+    {
+        return sessionProperties;
+    }
+
+    @JsonProperty
+    public Map> getCatalogSessionProperties()
+    {
+        return catalogSessionProperties;
+    }
+
+    public static class Mapper
+            implements RowMapper
+    {
+        @Override
+        public SessionMatchSpec map(ResultSet resultSet, StatementContext context)
+                throws SQLException
+        {
+            SessionInfo sessionInfo = getProperties(
+                    Optional.ofNullable(resultSet.getString("session_property_names")),
+                    Optional.ofNullable(resultSet.getString("session_property_values")),
+                    Optional.ofNullable(resultSet.getString("session_property_catalogs")));
+
+            return new SessionMatchSpec(
+                    Optional.ofNullable(resultSet.getString("user_regex")).map(Pattern::compile),
+                    Optional.ofNullable(resultSet.getString("source_regex")).map(Pattern::compile),
+                    Optional.ofNullable(resultSet.getString("client_tags")).map(tag -> Splitter.on(",").splitToList(tag)),
+                    Optional.ofNullable(resultSet.getString("query_type")),
+                    Optional.ofNullable(resultSet.getString("group_regex")).map(Pattern::compile),
+                    Optional.ofNullable(resultSet.getString("client_info_regex")).map(Pattern::compile),
+                    Optional.of(resultSet.getBoolean("override_session_properties")),
+                    sessionInfo.getSessionProperties(),
+                    sessionInfo.getCatalogProperties());
+        }
+
+        private SessionInfo getProperties(Optional names, Optional values, Optional catalogs)
+        {
+            if (!names.isPresent()) {
+                return new SessionInfo(ImmutableMap.of(), ImmutableMap.of());
+            }
+
+            checkArgument(catalogs.isPresent(), "names are present, but catalogs are not");
+            checkArgument(values.isPresent(), "names are present, but values are not");
+            List sessionPropertyNames = Splitter.on(",").splitToList(names.get());
+            List sessionPropertyValues = Splitter.on(",").splitToList(values.get());
+            List sessionPropertyCatalogs = Splitter.on(",").splitToList(catalogs.get());
+            checkArgument(sessionPropertyNames.size() == sessionPropertyValues.size(),
+                    "The number of property names and values should be the same");
+
+            checkArgument(sessionPropertyNames.size() == sessionPropertyCatalogs.size(),
+                    "The number of property names and catalog values should be the same");
+
+            Map> catalogSessionProperties = new HashMap<>();
+            Map systemSessionProperties = new HashMap<>();
+            for (int i = 0; i < sessionPropertyNames.size(); i++) {
+                if (sessionPropertyCatalogs.get(i).equals(emptyCatalog)) {
+                    systemSessionProperties.put(sessionPropertyNames.get(i), sessionPropertyValues.get(i));
+                }
+                else {
+                    catalogSessionProperties
+                            .computeIfAbsent(sessionPropertyCatalogs.get(i), k -> new HashMap<>())
+                            .put(sessionPropertyNames.get(i), sessionPropertyValues.get(i));
+                }
+            }
+
+            return new SessionInfo(systemSessionProperties, catalogSessionProperties);
+        }
+        static class SessionInfo
+        {
+            private final Map sessionProperties;
+            private final Map> catalogProperties;
+
+            public SessionInfo(Map sessionProperties,
+                               Map> catalogProperties)
+            {
+                this.sessionProperties = sessionProperties;
+                this.catalogProperties = catalogProperties;
+            }
+
+            public Map getSessionProperties()
+            {
+                return sessionProperties;
+            }
+
+            public Map> getCatalogProperties()
+            {
+                return catalogProperties;
+            }
+        }
+    }
+}
diff --git a/presto-session-property-managers/src/test/java/com/facebook/presto/session/TestFileSessionPropertyManager.java b/presto-session-property-managers-common/src/test/java/com/facebook/presto/session/AbstractTestSessionPropertyManager.java
similarity index 72%
rename from presto-session-property-managers/src/test/java/com/facebook/presto/session/TestFileSessionPropertyManager.java
rename to presto-session-property-managers-common/src/test/java/com/facebook/presto/session/AbstractTestSessionPropertyManager.java
index efd2c0f6b6411..3f09af61b1562 100644
--- a/presto-session-property-managers/src/test/java/com/facebook/presto/session/TestFileSessionPropertyManager.java
+++ b/presto-session-property-managers-common/src/test/java/com/facebook/presto/session/AbstractTestSessionPropertyManager.java
@@ -11,34 +11,25 @@
  * See the License for the specific language governing permissions and
  * limitations under the License.
  */
-
 package com.facebook.presto.session;
 
-import com.facebook.airlift.testing.TempFile;
 import com.facebook.presto.common.resourceGroups.QueryType;
 import com.facebook.presto.spi.resourceGroups.ResourceGroupId;
 import com.facebook.presto.spi.session.SessionConfigurationContext;
-import com.facebook.presto.spi.session.SessionPropertyConfigurationManager;
-import com.facebook.presto.spi.session.SessionPropertyConfigurationManager.SystemSessionPropertyConfiguration;
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableMap;
 import com.google.common.collect.ImmutableSet;
 import org.testng.annotations.Test;
 
 import java.io.IOException;
-import java.nio.file.Files;
-import java.nio.file.Path;
-import java.util.Arrays;
 import java.util.Map;
 import java.util.Optional;
 import java.util.regex.Pattern;
 
-import static com.facebook.presto.session.FileSessionPropertyManager.CODEC;
-import static org.testng.Assert.assertEquals;
-
-public class TestFileSessionPropertyManager
+@Test(singleThreaded = true)
+public abstract class AbstractTestSessionPropertyManager
 {
-    private static final SessionConfigurationContext CONTEXT = new SessionConfigurationContext(
+    protected static final SessionConfigurationContext CONTEXT = new SessionConfigurationContext(
             "user",
             Optional.empty(),
             Optional.of("source"),
@@ -48,6 +39,13 @@ public class TestFileSessionPropertyManager
             Optional.of("bar"),
             "testversion");
 
+    protected abstract void assertProperties(Map properties, SessionMatchSpec... spec)
+            throws IOException;
+    protected abstract void assertProperties(Map defaultProperties, Map overrideProperties, SessionMatchSpec... spec)
+            throws IOException;
+    protected abstract void assertProperties(Map defaultProperties, Map overrideProperties, Map> catalogProperties, SessionMatchSpec... spec)
+            throws IOException;
+
     @Test
     public void testResourceGroupMatch()
             throws IOException
@@ -61,7 +59,8 @@ public void testResourceGroupMatch()
                 Optional.of(Pattern.compile("global.pipeline.user_.*")),
                 Optional.empty(),
                 Optional.empty(),
-                properties);
+                properties,
+                ImmutableMap.of());
 
         assertProperties(properties, spec);
     }
@@ -79,7 +78,8 @@ public void testClientTagMatch()
                 Optional.empty(),
                 Optional.empty(),
                 Optional.empty(),
-                properties);
+                properties,
+                ImmutableMap.of());
 
         assertProperties(properties, spec);
     }
@@ -96,7 +96,8 @@ public void testMultipleMatch()
                 Optional.empty(),
                 Optional.empty(),
                 Optional.empty(),
-                ImmutableMap.of("PROPERTY1", "VALUE1"));
+                ImmutableMap.of("PROPERTY1", "VALUE1", "PROPERTY3", "VALUE3"),
+                ImmutableMap.of());
         SessionMatchSpec spec2 = new SessionMatchSpec(
                 Optional.empty(),
                 Optional.empty(),
@@ -105,9 +106,10 @@ public void testMultipleMatch()
                 Optional.empty(),
                 Optional.empty(),
                 Optional.empty(),
-                ImmutableMap.of("PROPERTY1", "VALUE1", "PROPERTY2", "VALUE2"));
+                ImmutableMap.of("PROPERTY1", "VALUE1", "PROPERTY2", "VALUE2"),
+                ImmutableMap.of());
 
-        assertProperties(ImmutableMap.of("PROPERTY1", "VALUE1", "PROPERTY2", "VALUE2"), spec1, spec2);
+        assertProperties(ImmutableMap.of("PROPERTY1", "VALUE1", "PROPERTY2", "VALUE2", "PROPERTY3", "VALUE3"), spec1, spec2);
     }
 
     @Test
@@ -122,7 +124,8 @@ public void testNoMatch()
                 Optional.of(Pattern.compile("global.interactive.user_.*")),
                 Optional.empty(),
                 Optional.empty(),
-                ImmutableMap.of("PROPERTY", "VALUE"));
+                ImmutableMap.of("PROPERTY", "VALUE"),
+                ImmutableMap.of());
 
         assertProperties(ImmutableMap.of(), spec);
     }
@@ -140,7 +143,8 @@ public void testClientInfoMatch()
                 Optional.empty(),
                 Optional.of(Pattern.compile("bar")),
                 Optional.empty(),
-                properties);
+                properties,
+                ImmutableMap.of());
 
         assertProperties(properties, spec);
     }
@@ -160,7 +164,8 @@ public void testOverride()
                 Optional.empty(),
                 Optional.of(Pattern.compile("bar")),
                 Optional.of(true),
-                overrideProperties);
+                overrideProperties,
+                ImmutableMap.of());
 
         SessionMatchSpec specDefault = new SessionMatchSpec(
                 Optional.empty(),
@@ -170,7 +175,8 @@ public void testOverride()
                 Optional.empty(),
                 Optional.of(Pattern.compile("bar")),
                 Optional.empty(),
-                defaultProperties);
+                defaultProperties,
+                ImmutableMap.of());
 
         // PROPERTY1 should be an override property with the value from the default (non-override, higher precedence)
         // spec.
@@ -178,22 +184,23 @@ public void testOverride()
         assertProperties(defaultProperties, ImmutableMap.of("PROPERTY1", "VALUE2"), specOverride, specDefault);
     }
 
-    private static void assertProperties(Map defaultProperties, SessionMatchSpec... spec)
+    @Test
+    public void testCatalogProperty()
             throws IOException
     {
-        assertProperties(defaultProperties, ImmutableMap.of(), spec);
-    }
+        ImmutableMap defaultProperties = ImmutableMap.of("PROPERTY1", "VALUE1");
+        ImmutableMap> catalogProperties = ImmutableMap.of("CATALOG", ImmutableMap.of("PROPERTY", "VALUE"));
+        SessionMatchSpec spec = new SessionMatchSpec(
+                Optional.empty(),
+                Optional.empty(),
+                Optional.empty(),
+                Optional.empty(),
+                Optional.empty(),
+                Optional.empty(),
+                Optional.empty(),
+                defaultProperties,
+                catalogProperties);
 
-    private static void assertProperties(Map defaultProperties, Map overrideProperties, SessionMatchSpec... spec)
-            throws IOException
-    {
-        try (TempFile tempFile = new TempFile()) {
-            Path configurationFile = tempFile.path();
-            Files.write(configurationFile, CODEC.toJsonBytes(Arrays.asList(spec)));
-            SessionPropertyConfigurationManager manager = new FileSessionPropertyManager(new FileSessionPropertyManagerConfig().setConfigFile(configurationFile.toFile()));
-            SystemSessionPropertyConfiguration propertyConfiguration = manager.getSystemSessionProperties(CONTEXT);
-            assertEquals(propertyConfiguration.systemPropertyDefaults, defaultProperties);
-            assertEquals(propertyConfiguration.systemPropertyOverrides, overrideProperties);
-        }
+        assertProperties(defaultProperties, ImmutableMap.of(), catalogProperties, spec);
     }
 }
diff --git a/presto-session-property-managers/src/main/java/com/facebook/presto/session/SessionMatchSpec.java b/presto-session-property-managers/src/main/java/com/facebook/presto/session/SessionMatchSpec.java
deleted file mode 100644
index 18f34a80adf9d..0000000000000
--- a/presto-session-property-managers/src/main/java/com/facebook/presto/session/SessionMatchSpec.java
+++ /dev/null
@@ -1,152 +0,0 @@
-/*
- * Licensed 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 com.facebook.presto.session;
-
-import com.facebook.presto.spi.resourceGroups.ResourceGroupId;
-import com.facebook.presto.spi.session.SessionConfigurationContext;
-import com.fasterxml.jackson.annotation.JsonCreator;
-import com.fasterxml.jackson.annotation.JsonProperty;
-import com.google.common.collect.ImmutableList;
-import com.google.common.collect.ImmutableMap;
-import com.google.common.collect.ImmutableSet;
-
-import java.util.List;
-import java.util.Map;
-import java.util.Optional;
-import java.util.Set;
-import java.util.regex.Pattern;
-
-import static java.util.Objects.requireNonNull;
-
-public class SessionMatchSpec
-{
-    private final Optional userRegex;
-    private final Optional sourceRegex;
-    private final Set clientTags;
-    private final Optional queryType;
-    private final Optional clientInfoRegex;
-    private final Optional resourceGroupRegex;
-    private final Optional overrideSessionProperties;
-    private final Map sessionProperties;
-
-    @JsonCreator
-    public SessionMatchSpec(
-            @JsonProperty("user") Optional userRegex,
-            @JsonProperty("source") Optional sourceRegex,
-            @JsonProperty("clientTags") Optional> clientTags,
-            @JsonProperty("queryType") Optional queryType,
-            @JsonProperty("group") Optional resourceGroupRegex,
-            @JsonProperty("clientInfo") Optional clientInfoRegex,
-            @JsonProperty("overrideSessionProperties") Optional overrideSessionProperties,
-            @JsonProperty("sessionProperties") Map sessionProperties)
-    {
-        this.userRegex = requireNonNull(userRegex, "userRegex is null");
-        this.sourceRegex = requireNonNull(sourceRegex, "sourceRegex is null");
-        requireNonNull(clientTags, "clientTags is null");
-        this.clientTags = ImmutableSet.copyOf(clientTags.orElse(ImmutableList.of()));
-        this.queryType = requireNonNull(queryType, "queryType is null");
-        this.resourceGroupRegex = requireNonNull(resourceGroupRegex, "resourceGroupRegex is null");
-        this.clientInfoRegex = requireNonNull(clientInfoRegex, "clientInfoRegex is null");
-        this.overrideSessionProperties = requireNonNull(overrideSessionProperties, "overrideSessionProperties is null");
-        requireNonNull(sessionProperties, "sessionProperties is null");
-        this.sessionProperties = ImmutableMap.copyOf(sessionProperties);
-    }
-
-    public Map match(SessionConfigurationContext context)
-    {
-        if (userRegex.isPresent() && !userRegex.get().matcher(context.getUser()).matches()) {
-            return ImmutableMap.of();
-        }
-        if (sourceRegex.isPresent()) {
-            String source = context.getSource().orElse("");
-            if (!sourceRegex.get().matcher(source).matches()) {
-                return ImmutableMap.of();
-            }
-        }
-        if (!clientTags.isEmpty() && !context.getClientTags().containsAll(clientTags)) {
-            return ImmutableMap.of();
-        }
-
-        if (queryType.isPresent()) {
-            String contextQueryType = context.getQueryType().orElse("");
-            if (!queryType.get().equalsIgnoreCase(contextQueryType)) {
-                return ImmutableMap.of();
-            }
-        }
-
-        if (clientInfoRegex.isPresent()) {
-            String clientInfo = context.getClientInfo().orElse("");
-            if (!clientInfoRegex.get().matcher(clientInfo).matches()) {
-                return ImmutableMap.of();
-            }
-        }
-
-        if (resourceGroupRegex.isPresent()) {
-            String resourceGroupId = context.getResourceGroupId().map(ResourceGroupId::toString).orElse("");
-            if (!resourceGroupRegex.get().matcher(resourceGroupId).matches()) {
-                return ImmutableMap.of();
-            }
-        }
-
-        return sessionProperties;
-    }
-
-    @JsonProperty("user")
-    public Optional getUserRegex()
-    {
-        return userRegex;
-    }
-
-    @JsonProperty("source")
-    public Optional getSourceRegex()
-    {
-        return sourceRegex;
-    }
-
-    @JsonProperty
-    public Set getClientTags()
-    {
-        return clientTags;
-    }
-
-    @JsonProperty
-    public Optional getQueryType()
-    {
-        return queryType;
-    }
-
-    @JsonProperty("group")
-    public Optional getResourceGroupRegex()
-    {
-        return resourceGroupRegex;
-    }
-
-    @JsonProperty
-    public Optional getClientInfoRegex()
-    {
-        return clientInfoRegex;
-    }
-
-    @JsonProperty
-    public Optional getOverrideSessionProperties()
-    {
-        return overrideSessionProperties;
-    }
-
-    @JsonProperty
-    public Map getSessionProperties()
-    {
-        return sessionProperties;
-    }
-}
diff --git a/presto-singlestore/src/test/java/com/facebook/presto/plugin/singlestore/TestSingleStoreDistributedQueries.java b/presto-singlestore/src/test/java/com/facebook/presto/plugin/singlestore/TestSingleStoreDistributedQueries.java
index 7a69686ebd915..dbe8acb45a5bb 100644
--- a/presto-singlestore/src/test/java/com/facebook/presto/plugin/singlestore/TestSingleStoreDistributedQueries.java
+++ b/presto-singlestore/src/test/java/com/facebook/presto/plugin/singlestore/TestSingleStoreDistributedQueries.java
@@ -183,6 +183,18 @@ public void testDescribeOutputNamedAndUnnamed()
         // this connector uses a non-canonical type for varchar columns in tpch
     }
 
+    @Override
+    public void testNonAutoCommitTransactionWithRollback()
+    {
+        // Catalog singlestore only supports writes using autocommit
+    }
+
+    @Override
+    public void testNonAutoCommitTransactionWithCommit()
+    {
+        // Catalog singlestore only supports writes using autocommit
+    }
+
     @Override
     public void testInsert()
     {
diff --git a/presto-spark-base/pom.xml b/presto-spark-base/pom.xml
index 8f3dd0f91a66d..7767219a7ec65 100644
--- a/presto-spark-base/pom.xml
+++ b/presto-spark-base/pom.xml
@@ -16,7 +16,6 @@
         9.4.55.v20240627
         4.12.0
         3.9.1
-        2
     
 
     
@@ -55,12 +54,6 @@
             provided
         
 
-        
-            com.facebook.presto
-            presto-spark-classloader-spark${dep.pos.classloader.module-name.suffix}
-            provided
-        
-
         
             com.facebook.presto
             presto-client
@@ -117,6 +110,10 @@
                     com.facebook.airlift.drift
                     *
                 
+                
+                    com.facebook.presto
+                    presto-function-namespace-managers-common
+                
             
         
 
@@ -534,6 +531,25 @@
                 
             
         
+        
+            spark2
+
+            
+                true
+                
+                    !spark-version
+                
+            
+
+            
+                
+                    com.facebook.presto
+                    presto-spark-classloader-spark2
+                    ${project.version}
+                
+            
+        
+
         
             spark3
 
@@ -544,11 +560,13 @@
                 
             
 
-            
-                3
-            
-
             
+                
+                    com.facebook.presto
+                    presto-spark-classloader-spark3
+                    ${project.version}
+                
+
                 
                     com.facebook.presto.spark
                     spark-core
diff --git a/presto-spark-base/src/main/java/com/facebook/presto/spark/PrestoSparkInjectorFactory.java b/presto-spark-base/src/main/java/com/facebook/presto/spark/PrestoSparkInjectorFactory.java
index 9608adf93f843..bfabba20b632e 100644
--- a/presto-spark-base/src/main/java/com/facebook/presto/spark/PrestoSparkInjectorFactory.java
+++ b/presto-spark-base/src/main/java/com/facebook/presto/spark/PrestoSparkInjectorFactory.java
@@ -30,6 +30,7 @@
 import com.facebook.presto.server.security.PrestoAuthenticatorManager;
 import com.facebook.presto.spark.classloader_interface.PrestoSparkBootstrapTimer;
 import com.facebook.presto.spark.classloader_interface.SparkProcessType;
+import com.facebook.presto.spark.execution.property.NativeExecutionConfigModule;
 import com.facebook.presto.spi.security.AccessControl;
 import com.facebook.presto.sql.analyzer.FeaturesConfig;
 import com.facebook.presto.sql.parser.SqlParserOptions;
@@ -165,6 +166,11 @@ public Injector create(PrestoSparkBootstrapTimer bootstrapTimer)
             modules.add(new TempStorageModule());
         }
 
+        Map nativeWorkerConfigs = new HashMap<>(
+                this.nativeWorkerConfigProperties.orElse(ImmutableMap.of()));
+        nativeWorkerConfigs.put("node.environment", "spark");
+        modules.add(new NativeExecutionConfigModule(nativeWorkerConfigs));
+
         modules.addAll(additionalModules);
 
         Bootstrap app = new Bootstrap(modules.build());
@@ -172,11 +178,10 @@ public Injector create(PrestoSparkBootstrapTimer bootstrapTimer)
         // Stream redirect doesn't work well with spark logging
         app.doNotInitializeLogging();
 
-        Map requiredProperties = new HashMap<>();
-        requiredProperties.put("node.environment", "spark");
-        requiredProperties.putAll(configProperties);
-
-        app.setRequiredConfigurationProperties(ImmutableMap.copyOf(requiredProperties));
+        Map requiredConfigProperties = new HashMap<>();
+        requiredConfigProperties.put("node.environment", "spark");
+        requiredConfigProperties.putAll(configProperties);
+        app.setRequiredConfigurationProperties(ImmutableMap.copyOf(requiredConfigProperties));
 
         bootstrapTimer.beginInjectorInitialization();
         Injector injector = app.initialize();
diff --git a/presto-spark-base/src/main/java/com/facebook/presto/spark/PrestoSparkModule.java b/presto-spark-base/src/main/java/com/facebook/presto/spark/PrestoSparkModule.java
index 650d533492476..5cd272b85b277 100644
--- a/presto-spark-base/src/main/java/com/facebook/presto/spark/PrestoSparkModule.java
+++ b/presto-spark-base/src/main/java/com/facebook/presto/spark/PrestoSparkModule.java
@@ -61,6 +61,7 @@
 import com.facebook.presto.execution.resourceGroups.InternalResourceGroupManager;
 import com.facebook.presto.execution.resourceGroups.ResourceGroupManager;
 import com.facebook.presto.execution.scheduler.NodeSchedulerConfig;
+import com.facebook.presto.execution.scheduler.clusterOverload.ClusterOverloadPolicyModule;
 import com.facebook.presto.execution.scheduler.nodeSelection.SimpleTtlNodeSelectorConfig;
 import com.facebook.presto.execution.warnings.WarningCollectorConfig;
 import com.facebook.presto.index.IndexManager;
@@ -120,9 +121,7 @@
 import com.facebook.presto.spark.execution.PrestoSparkBroadcastTableCacheManager;
 import com.facebook.presto.spark.execution.PrestoSparkExecutionExceptionFactory;
 import com.facebook.presto.spark.execution.http.BatchTaskUpdateRequest;
-import com.facebook.presto.spark.execution.property.NativeExecutionConnectorConfig;
 import com.facebook.presto.spark.execution.property.NativeExecutionNodeConfig;
-import com.facebook.presto.spark.execution.property.NativeExecutionSystemConfig;
 import com.facebook.presto.spark.execution.shuffle.PrestoSparkLocalShuffleReadInfo;
 import com.facebook.presto.spark.execution.shuffle.PrestoSparkLocalShuffleWriteInfo;
 import com.facebook.presto.spark.execution.task.PrestoSparkNativeTaskExecutorFactory;
@@ -284,9 +283,7 @@ protected void setup(Binder binder)
         configBinder(binder).bindConfig(SessionPropertyProviderConfig.class);
         configBinder(binder).bindConfig(PrestoSparkConfig.class);
         configBinder(binder).bindConfig(TracingConfig.class);
-        configBinder(binder).bindConfig(NativeExecutionSystemConfig.class);
         configBinder(binder).bindConfig(NativeExecutionNodeConfig.class);
-        configBinder(binder).bindConfig(NativeExecutionConnectorConfig.class);
         configBinder(binder).bindConfig(PlanCheckerProviderManagerConfig.class);
         configBinder(binder).bindConfig(SecurityConfig.class);
 
@@ -334,6 +331,9 @@ protected void setup(Binder binder)
         // handle resolver
         binder.install(new HandleJsonModule());
 
+        // ClusterOverload policy module
+        binder.install(new ClusterOverloadPolicyModule());
+
         // plugin manager
         configBinder(binder).bindConfig(PluginManagerConfig.class);
         binder.bind(PluginManager.class).in(Scopes.SINGLETON);
diff --git a/presto-spark-base/src/main/java/com/facebook/presto/spark/PrestoSparkQueryExecutionFactory.java b/presto-spark-base/src/main/java/com/facebook/presto/spark/PrestoSparkQueryExecutionFactory.java
index 042b56ce62e1a..8eef5d1db658b 100644
--- a/presto-spark-base/src/main/java/com/facebook/presto/spark/PrestoSparkQueryExecutionFactory.java
+++ b/presto-spark-base/src/main/java/com/facebook/presto/spark/PrestoSparkQueryExecutionFactory.java
@@ -679,7 +679,7 @@ else if (preparedQuery.isExplainTypeValidate()) {
                 planAndMore = queryPlanner.createQueryPlan(session, preparedQuery, warningCollector, variableAllocator, planNodeIdAllocator, sparkContext, sql);
                 JavaSparkContext javaSparkContext = new JavaSparkContext(sparkContext);
                 CollectionAccumulator taskInfoCollector = new CollectionAccumulator<>();
-                taskInfoCollector.register(sparkContext, Option.empty(), false);
+                taskInfoCollector.register(sparkContext, Option.empty(), true);
                 CollectionAccumulator shuffleStatsCollector = new CollectionAccumulator<>();
                 shuffleStatsCollector.register(sparkContext, Option.empty(), false);
                 TempStorage tempStorage = tempStorageManager.getTempStorage(storageBasedBroadcastJoinStorage);
diff --git a/presto-spark-base/src/main/java/com/facebook/presto/spark/PrestoSparkServiceFactory.java b/presto-spark-base/src/main/java/com/facebook/presto/spark/PrestoSparkServiceFactory.java
index 741b018631c62..c86d3b30f94ef 100644
--- a/presto-spark-base/src/main/java/com/facebook/presto/spark/PrestoSparkServiceFactory.java
+++ b/presto-spark-base/src/main/java/com/facebook/presto/spark/PrestoSparkServiceFactory.java
@@ -27,6 +27,7 @@
 import com.google.inject.Module;
 
 import java.util.List;
+import java.util.Optional;
 
 import static com.facebook.presto.spark.classloader_interface.PrestoSparkConfiguration.METADATA_STORAGE_TYPE_LOCAL;
 import static com.google.common.base.Preconditions.checkArgument;
@@ -70,7 +71,7 @@ protected List getAdditionalModules(PrestoSparkConfiguration configurati
         return ImmutableList.of(
                 new PrestoSparkLocalMetadataStorageModule(),
                 // TODO: Need to let NativeExecutionModule addition be controlled by configuration as well.
-                new NativeExecutionModule());
+                new NativeExecutionModule(Optional.of(configuration.getCatalogProperties())));
     }
 
     protected SqlParserOptions getSqlParserOptions()
diff --git a/presto-spark-base/src/main/java/com/facebook/presto/spark/execution/AbstractPrestoSparkQueryExecution.java b/presto-spark-base/src/main/java/com/facebook/presto/spark/execution/AbstractPrestoSparkQueryExecution.java
index 2e03dd64a54d1..b76da7f2ddc2e 100644
--- a/presto-spark-base/src/main/java/com/facebook/presto/spark/execution/AbstractPrestoSparkQueryExecution.java
+++ b/presto-spark-base/src/main/java/com/facebook/presto/spark/execution/AbstractPrestoSparkQueryExecution.java
@@ -29,7 +29,9 @@
 import com.facebook.presto.execution.QueryState;
 import com.facebook.presto.execution.QueryStateTimer;
 import com.facebook.presto.execution.StageInfo;
+import com.facebook.presto.execution.TaskId;
 import com.facebook.presto.execution.TaskInfo;
+import com.facebook.presto.execution.TaskState;
 import com.facebook.presto.execution.scheduler.ExecutionWriterTarget;
 import com.facebook.presto.execution.scheduler.StreamingPlanSection;
 import com.facebook.presto.execution.scheduler.StreamingSubPlan;
@@ -116,6 +118,7 @@
 import java.util.TreeMap;
 import java.util.concurrent.TimeoutException;
 import java.util.concurrent.atomic.AtomicReference;
+import java.util.stream.Collectors;
 import java.util.stream.IntStream;
 
 import static com.facebook.airlift.units.DataSize.Unit.BYTE;
@@ -572,23 +575,78 @@ protected void validateStorageCapabilities(TempStorage tempStorage)
         }
     }
 
+    /**
+     * Updates the taskInfoMap to ensure it stores the most relevant {@link TaskInfo} for each
+     * logical task, identified by task ID (excluding attempt number).
+     * 

+     * This method ensures that, for each logical task, the map retains the latest successful
+     * attempt if available, or otherwise the most recent attempt based on attempt number. Warnings
+     * are logged in cases of unexpected duplicate or multiple successful attempts.
+     *
+     * @param taskInfoMap the map from logical task ID (taskId excluding attempt number) to
+     * {@link TaskInfo}
+     * @param taskInfo the {@link TaskInfo} to consider for updating the map
+     */
+    private void updateTaskInfoMap(HashMap taskInfoMap, TaskInfo taskInfo)
+    {
+        TaskId newTaskId = taskInfo.getTaskId();
+        String taskIdWithoutAttemptId = new StringBuilder()
+                .append(newTaskId.getStageExecutionId().toString())
+                .append(".")
+                .append(newTaskId.getId())
+                .toString();
+        if (!taskInfoMap.containsKey(taskIdWithoutAttemptId)) {
+            taskInfoMap.put(taskIdWithoutAttemptId, taskInfo);
+            return;
+        }
+
+        TaskInfo storedTaskInfo = taskInfoMap.get(taskIdWithoutAttemptId);
+        TaskId storedTaskId = storedTaskInfo.getTaskId();
+        TaskState storedTaskState = storedTaskInfo.getTaskStatus().getState();
+        TaskState newTaskState = taskInfo.getTaskStatus().getState();
+        if (storedTaskState == TaskState.FINISHED) {
+            if (newTaskState == TaskState.FINISHED) {
+                log.warn("Multiple attempts of the same task have succeeded %s vs %s",
+                        storedTaskId.toString(), newTaskId.toString());
+            }
+            // Successful one has been stored. Nothing needs to be done.
+            return;
+        }
+
+        int storedAttemptNumber = storedTaskId.getAttemptNumber();
+        int newAttemptNumber = newTaskId.getAttemptNumber();
+        if (newTaskState == TaskState.FINISHED || storedAttemptNumber < newAttemptNumber) {
+            taskInfoMap.put(taskIdWithoutAttemptId, taskInfo);
+        }
+        if (storedAttemptNumber == newAttemptNumber) {
+            log.warn("Received multiple identical TaskId %s vs %s",
+                    storedTaskId.toString(), newTaskId.toString());
+        }
+    }
+
     protected void queryCompletedEvent(Optional failureInfo, OptionalLong updateCount)
     {
         List serializedTaskInfos = taskInfoCollector.value();
-        ImmutableList.Builder taskInfos = ImmutableList.builder();
+        HashMap taskInfoMap = new HashMap<>();
         long totalSerializedTaskInfoSizeInBytes = 0;
         for (SerializedTaskInfo serializedTaskInfo : serializedTaskInfos) {
             byte[] bytes = serializedTaskInfo.getBytesAndClear();
             totalSerializedTaskInfoSizeInBytes += bytes.length;
             TaskInfo taskInfo = deserializeZstdCompressed(taskInfoCodec, bytes);
-            taskInfos.add(taskInfo);
+            updateTaskInfoMap(taskInfoMap, taskInfo);
         }
         taskInfoCollector.reset();
 
-        log.info("Total serialized task info size: %s", DataSize.succinctBytes(totalSerializedTaskInfoSizeInBytes));
+        log.info("Total serialized task info count %s size: %s. Total deduped task info count %s",
+                serializedTaskInfos.size(),
+                DataSize.succinctBytes(totalSerializedTaskInfoSizeInBytes),
+                taskInfoMap.size());
 
         Optional stageInfoOptional = getFinalFragmentedPlan().map(finalFragmentedPlan ->
-                PrestoSparkQueryExecutionFactory.createStageInfo(session.getQueryId(), finalFragmentedPlan, taskInfos.build()));
+                PrestoSparkQueryExecutionFactory.createStageInfo(
+                        session.getQueryId(),
+                        finalFragmentedPlan,
+                        taskInfoMap.values().stream().collect(Collectors.toList())));
         QueryState queryState = failureInfo.isPresent() ? FAILED : FINISHED;
 
         QueryInfo queryInfo = PrestoSparkQueryExecutionFactory.createQueryInfo(
diff --git a/presto-spark-base/src/main/java/com/facebook/presto/spark/execution/nativeprocess/NativeExecutionModule.java b/presto-spark-base/src/main/java/com/facebook/presto/spark/execution/nativeprocess/NativeExecutionModule.java
index d265938edf999..8df6a83395b35 100644
--- a/presto-spark-base/src/main/java/com/facebook/presto/spark/execution/nativeprocess/NativeExecutionModule.java
+++ b/presto-spark-base/src/main/java/com/facebook/presto/spark/execution/nativeprocess/NativeExecutionModule.java
@@ -13,9 +13,7 @@
  */
 package com.facebook.presto.spark.execution.nativeprocess;
 
-import com.facebook.presto.spark.execution.property.NativeExecutionConnectorConfig;
-import com.facebook.presto.spark.execution.property.NativeExecutionNodeConfig;
-import com.facebook.presto.spark.execution.property.NativeExecutionSystemConfig;
+import com.facebook.presto.spark.execution.property.NativeExecutionCatalogProperties;
 import com.facebook.presto.spark.execution.property.PrestoSparkWorkerProperty;
 import com.facebook.presto.spark.execution.property.WorkerProperty;
 import com.facebook.presto.spark.execution.shuffle.PrestoSparkLocalShuffleInfoTranslator;
@@ -23,6 +21,7 @@
 import com.facebook.presto.spark.execution.task.ForNativeExecutionTask;
 import com.facebook.presto.spark.execution.task.NativeExecutionTaskFactory;
 import com.google.common.annotations.VisibleForTesting;
+import com.google.common.collect.ImmutableMap;
 import com.google.inject.Binder;
 import com.google.inject.Module;
 import com.google.inject.Provides;
@@ -31,6 +30,7 @@
 import com.google.inject.TypeLiteral;
 import okhttp3.OkHttpClient;
 
+import java.util.Map;
 import java.util.Optional;
 import java.util.concurrent.TimeUnit;
 
@@ -39,20 +39,14 @@
 public class NativeExecutionModule
         implements Module
 {
-    private Optional connectorConfig;
-
-    // For use by production system where the configurations can only be tuned via configurations.
-    public NativeExecutionModule()
-    {
-        this.connectorConfig = Optional.empty();
-    }
+    private Optional>> catalogProperties;
 
     // In the future, we would make more bindings injected into NativeExecutionModule
     // to be able to test various configuration parameters
     @VisibleForTesting
-    public NativeExecutionModule(Optional connectorConfig)
+    public NativeExecutionModule(Optional>> catalogProperties)
     {
-        this.connectorConfig = connectorConfig;
+        this.catalogProperties = catalogProperties;
     }
 
     @Override
@@ -73,13 +67,14 @@ protected void bindShuffle(Binder binder)
 
     protected void bindWorkerProperties(Binder binder)
     {
-        newOptionalBinder(binder, new TypeLiteral>() {}).setDefault().to(PrestoSparkWorkerProperty.class).in(Scopes.SINGLETON);
-        if (connectorConfig.isPresent()) {
-            binder.bind(PrestoSparkWorkerProperty.class).toInstance(new PrestoSparkWorkerProperty(connectorConfig.get(), new NativeExecutionNodeConfig(), new NativeExecutionSystemConfig()));
-        }
-        else {
-            binder.bind(PrestoSparkWorkerProperty.class).in(Scopes.SINGLETON);
-        }
+        // Bind NativeExecutionCatalogProperties - this is not bound elsewhere
+        binder.bind(NativeExecutionCatalogProperties.class).toInstance(
+                new NativeExecutionCatalogProperties(catalogProperties.orElse(ImmutableMap.of())));
+
+        // Bind worker property classes
+        newOptionalBinder(binder, new TypeLiteral>() {
+        }).setDefault().to(PrestoSparkWorkerProperty.class).in(Scopes.SINGLETON);
+        binder.bind(PrestoSparkWorkerProperty.class).in(Scopes.SINGLETON);
     }
 
     protected void bindHttpClient(Binder binder)
diff --git a/presto-spark-base/src/main/java/com/facebook/presto/spark/execution/nativeprocess/NativeExecutionProcess.java b/presto-spark-base/src/main/java/com/facebook/presto/spark/execution/nativeprocess/NativeExecutionProcess.java
index b1d104cdbe9a7..c2c0ded8416a5 100644
--- a/presto-spark-base/src/main/java/com/facebook/presto/spark/execution/nativeprocess/NativeExecutionProcess.java
+++ b/presto-spark-base/src/main/java/com/facebook/presto/spark/execution/nativeprocess/NativeExecutionProcess.java
@@ -15,6 +15,7 @@
 
 import com.facebook.airlift.json.JsonCodec;
 import com.facebook.airlift.log.Logger;
+import com.facebook.airlift.units.DataSize;
 import com.facebook.airlift.units.Duration;
 import com.facebook.presto.Session;
 import com.facebook.presto.client.ServerInfo;
@@ -22,6 +23,8 @@
 import com.facebook.presto.spark.execution.http.PrestoSparkHttpServerClient;
 import com.facebook.presto.spark.execution.http.server.RequestErrorTracker;
 import com.facebook.presto.spark.execution.http.server.smile.BaseResponse;
+import com.facebook.presto.spark.execution.property.NativeExecutionSystemConfig;
+import com.facebook.presto.spark.execution.property.PrestoSparkWorkerProperty;
 import com.facebook.presto.spark.execution.property.WorkerProperty;
 import com.facebook.presto.spi.PrestoException;
 import com.google.common.annotations.VisibleForTesting;
@@ -31,6 +34,7 @@
 import com.google.common.util.concurrent.SettableFuture;
 import okhttp3.HttpUrl;
 import okhttp3.OkHttpClient;
+import org.apache.spark.SparkConf;
 import org.apache.spark.SparkEnv$;
 import org.apache.spark.SparkFiles;
 
@@ -60,11 +64,12 @@
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.atomic.AtomicBoolean;
 
+import static com.facebook.airlift.units.DataSize.Unit.BYTE;
+import static com.facebook.airlift.units.DataSize.Unit.GIGABYTE;
 import static com.facebook.presto.spi.StandardErrorCode.GENERIC_INTERNAL_ERROR;
 import static com.facebook.presto.spi.StandardErrorCode.NATIVE_EXECUTION_BINARY_NOT_EXIST;
 import static com.facebook.presto.spi.StandardErrorCode.NATIVE_EXECUTION_PROCESS_LAUNCH_ERROR;
 import static com.facebook.presto.spi.StandardErrorCode.NATIVE_EXECUTION_TASK_ERROR;
-import static com.google.common.base.Preconditions.checkArgument;
 import static com.google.common.util.concurrent.Futures.addCallback;
 import static com.google.common.util.concurrent.MoreExecutors.directExecutor;
 import static java.lang.String.format;
@@ -81,6 +86,7 @@ public class NativeExecutionProcess
     private static final String WORKER_CONFIG_FILE = "/config.properties";
     private static final String WORKER_NODE_CONFIG_FILE = "/node.properties";
     private static final String WORKER_CONNECTOR_CONFIG_FILE = "/catalog/";
+    private static final String NATIVE_PROCESS_MEMORY_SPARK_CONF_NAME = "spark.memory.offHeap.size";
     private static final int SIGSYS = 31;
 
     private final String executablePath;
@@ -130,6 +136,8 @@ public NativeExecutionProcess(
                 scheduledExecutorService,
                 "getting native process status");
         this.workerProperty = requireNonNull(workerProperty, "workerProperty is null");
+        // Update any runtime configs to be used by presto native worker
+        updateWorkerProperties();
     }
 
     /**
@@ -324,24 +332,111 @@ private static int getAvailableTcpPort(String nodeInternalAddress)
         }
     }
 
-    private String getNativeExecutionCatalogName(Session session)
+    private void populateConfigurationFiles(String configBasePath)
+            throws IOException
     {
-        checkArgument(session.getCatalog().isPresent(), "Catalog isn't set in the session.");
-        return session.getCatalog().get();
+        workerProperty.populateAllProperties(
+                Paths.get(configBasePath, WORKER_CONFIG_FILE),
+                Paths.get(configBasePath, WORKER_NODE_CONFIG_FILE),
+                Paths.get(configBasePath, WORKER_CONNECTOR_CONFIG_FILE));  // Directory path for catalogs
     }
 
-    private void populateConfigurationFiles(String configBasePath)
-            throws IOException
+    private void updateWorkerProperties()
     {
+        // Update memory properties
+        updateWorkerMemoryProperties();
+
         // The reason we have to pick and assign the port per worker is in our prod environment,
         // there is no port isolation among all the containers running on the same host, so we have
         // to pick unique port per worker to avoid port collision. This config will be passed down to
         // the native execution process eventually for process initialization.
-        workerProperty.getSystemConfig().setHttpServerPort(port);
-        workerProperty.populateAllProperties(
-                Paths.get(configBasePath, WORKER_CONFIG_FILE),
-                Paths.get(configBasePath, WORKER_NODE_CONFIG_FILE),
-                Paths.get(configBasePath, format("%s%s.properties", WORKER_CONNECTOR_CONFIG_FILE, getNativeExecutionCatalogName(session))));
+        workerProperty.getSystemConfig()
+                .update(NativeExecutionSystemConfig.HTTP_SERVER_HTTP_PORT, String.valueOf(port));
+    }
+
+    protected SparkConf getSparkConf()
+    {
+        return SparkEnv$.MODULE$.get() == null ? null : SparkEnv$.MODULE$.get().conf();
+    }
+
+    protected PrestoSparkWorkerProperty getWorkerProperty()
+    {
+        return (PrestoSparkWorkerProperty) workerProperty;
+    }
+
+    /**
+     * Computes values for system-memory-gb and query-memory-gb to start the native worker
+     * with.
+     * This logic is mainly useful when spark has provisioned larger containers to run
+     * previously OOMing tasks. Spark will provision larger container but without below
+     * logic the cpp process will not be able to use it.
+     *
+     * Also, we write the logic in a way that same logic applies during first attempt v/s
+     * subsequent OOMed larger container retry attempts
+     *
+     * The logic is simple and is as below
+     * - New system-memory-gb = spark.memory.offHeap.size
+     * - Then to calculate the new value of query-memory-gb we assume that
+     *   the new query-memory to system-memory ratio should be same as old values.
+     *   So we set newQueryMemory = newSystemMemory = (oldQueryMemory/oldSystemMemory)
+     *
+     *   TODO: In future make this algorithm more configurable. i.e. we might want a min/max
+     *         cap on the systemMemoryGb-queryMemoryGb buffer. Currently we just assume ratio
+     *         is good enough
+     */
+    protected void updateWorkerMemoryProperties()
+    {
+        // If sparkConf.NATIVE_PROCESS_MEMORY_SPARK_CONF_NAME is not set
+        // skip making any updates
+        SparkConf conf = getSparkConf();
+        if (conf == null) {
+            log.info("Not adjusting native process memory as conf is null");
+            return;
+        }
+        if (!conf.contains(NATIVE_PROCESS_MEMORY_SPARK_CONF_NAME)) {
+            log.info("Not adjusting native process memory as %s is not set", NATIVE_PROCESS_MEMORY_SPARK_CONF_NAME);
+            return;
+        }
+        DataSize offHeapMemoryBytes = DataSize.succinctDataSize(
+                conf.getSizeAsBytes(NATIVE_PROCESS_MEMORY_SPARK_CONF_NAME), BYTE);
+        DataSize currentSystemMemory = DataSize.valueOf(workerProperty.getSystemConfig().getAllProperties()
+                .get(NativeExecutionSystemConfig.SYSTEM_MEMORY_GB) + GIGABYTE.getUnitString());
+        DataSize currentQueryMemory = DataSize.valueOf(workerProperty.getSystemConfig().getAllProperties()
+                .get(NativeExecutionSystemConfig.QUERY_MEMORY_GB) + GIGABYTE.getUnitString());
+        if (offHeapMemoryBytes.toBytes() == 0
+                || currentSystemMemory.toBytes() == 0
+                || offHeapMemoryBytes.toBytes() < currentSystemMemory.toBytes()) {
+            log.info("Not adjusting native process memory as" +
+                    " offHeapMemoryBytes=%s,currentSystemMemory=%s are invalid", offHeapMemoryBytes, currentSystemMemory.toBytes());
+            return;
+        }
+
+        log.info("Setting Native Worker system-memory-gb to offHeap: %s", offHeapMemoryBytes);
+        DataSize newSystemMemory = offHeapMemoryBytes.convertTo(GIGABYTE);
+
+        double queryMemoryFraction = currentQueryMemory.toBytes() * 1.0 / currentSystemMemory.toBytes();
+        DataSize newQueryMemoryBytes = DataSize.succinctDataSize(
+                queryMemoryFraction * newSystemMemory.toBytes(), BYTE);
+        log.info("Dynamically Tuning Presto Native Memory Configs. " +
+                        "Configured SparkOffHeap: %s; " +
+                        "[oldSystemMemory: %s, newSystemMemory: %s], queryMemoryFraction: %s, " +
+                        "[oldQueryMemory: %s, newQueryMemory: %s]",
+                offHeapMemoryBytes,
+                currentSystemMemory,
+                newSystemMemory,
+                queryMemoryFraction,
+                currentQueryMemory,
+                newQueryMemoryBytes);
+
+        workerProperty.getSystemConfig()
+                .update(NativeExecutionSystemConfig.SYSTEM_MEMORY_GB,
+                        String.valueOf((int) newSystemMemory.getValue(GIGABYTE)));
+        workerProperty.getSystemConfig()
+                .update(NativeExecutionSystemConfig.QUERY_MEMORY_GB,
+                        String.valueOf((int) newQueryMemoryBytes.getValue(GIGABYTE)));
+        workerProperty.getSystemConfig()
+                .update(NativeExecutionSystemConfig.QUERY_MAX_MEMORY_PER_NODE,
+                        newQueryMemoryBytes.convertTo(GIGABYTE).toString());
     }
 
     private void doGetServerInfo(SettableFuture future)
diff --git a/presto-spark-base/src/main/java/com/facebook/presto/spark/execution/property/NativeExecutionCatalogProperties.java b/presto-spark-base/src/main/java/com/facebook/presto/spark/execution/property/NativeExecutionCatalogProperties.java
new file mode 100644
index 0000000000000..a217c4e7f5ae8
--- /dev/null
+++ b/presto-spark-base/src/main/java/com/facebook/presto/spark/execution/property/NativeExecutionCatalogProperties.java
@@ -0,0 +1,37 @@
+/*
+ * Licensed 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 com.facebook.presto.spark.execution.property;
+
+import java.util.Map;
+
+import static java.util.Objects.requireNonNull;
+
+/**
+ * This class holds catalog properties for native execution process.
+ * Each catalog will generate a separate .properties file.
+ */
+public class NativeExecutionCatalogProperties
+{
+    private final Map> catalogProperties;
+
+    public NativeExecutionCatalogProperties(Map> catalogProperties)
+    {
+        this.catalogProperties = requireNonNull(catalogProperties, "catalogProperties is null");
+    }
+
+    public Map> getAllCatalogProperties()
+    {
+        return catalogProperties;
+    }
+}
diff --git a/presto-spark-base/src/main/java/com/facebook/presto/spark/execution/property/NativeExecutionConfigModule.java b/presto-spark-base/src/main/java/com/facebook/presto/spark/execution/property/NativeExecutionConfigModule.java
new file mode 100644
index 0000000000000..38f6b365bdf79
--- /dev/null
+++ b/presto-spark-base/src/main/java/com/facebook/presto/spark/execution/property/NativeExecutionConfigModule.java
@@ -0,0 +1,51 @@
+/*
+ * Licensed 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 com.facebook.presto.spark.execution.property;
+
+import com.google.common.collect.ImmutableMap;
+import com.google.inject.AbstractModule;
+import com.google.inject.Scopes;
+import com.google.inject.TypeLiteral;
+import com.google.inject.name.Names;
+
+import java.util.Map;
+
+import static java.util.Objects.requireNonNull;
+
+/**
+ * Native configuration module that allows the population of config properties without config
+ * annotation checks.
+ */
+public class NativeExecutionConfigModule
+        extends AbstractModule
+{
+    private final Map systemConfigs;
+
+    public NativeExecutionConfigModule(Map systemConfigs)
+    {
+        this.systemConfigs = ImmutableMap.copyOf(
+                requireNonNull(systemConfigs, "systemConfigs is null"));
+    }
+
+    @Override
+    protected void configure()
+    {
+        bind(new TypeLiteral>() {})
+                .annotatedWith(
+                        Names.named(NativeExecutionSystemConfig.NATIVE_EXECUTION_SYSTEM_CONFIG))
+                .toInstance(systemConfigs);
+
+        bind(NativeExecutionSystemConfig.class).in(Scopes.SINGLETON);
+    }
+}
diff --git a/presto-spark-base/src/main/java/com/facebook/presto/spark/execution/property/NativeExecutionConnectorConfig.java b/presto-spark-base/src/main/java/com/facebook/presto/spark/execution/property/NativeExecutionConnectorConfig.java
deleted file mode 100644
index 34fc929b6da79..0000000000000
--- a/presto-spark-base/src/main/java/com/facebook/presto/spark/execution/property/NativeExecutionConnectorConfig.java
+++ /dev/null
@@ -1,86 +0,0 @@
-/*
- * Licensed 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 com.facebook.presto.spark.execution.property;
-
-import com.facebook.airlift.configuration.Config;
-import com.facebook.airlift.units.DataSize;
-import com.google.common.collect.ImmutableMap;
-
-import java.util.Map;
-
-/**
- * This config class corresponds to catalog/.properties for native execution process. Properties inside will be used in PrestoServer.cpp
- * Currently, presto-on-spark native only supports querying 1 catalog in a single spark session.
- */
-public class NativeExecutionConnectorConfig
-{
-    private static final String CACHE_ENABLED = "cache.enabled";
-    private static final String CACHE_MAX_CACHE_SIZE = "cache.max-cache-size";
-    private static final String CONNECTOR_NAME = "connector.name";
-
-    private boolean cacheEnabled;
-    // maxCacheSize in MB
-    private DataSize maxCacheSize = new DataSize(0, DataSize.Unit.MEGABYTE);
-    private String connectorName = "hive";
-
-    public Map getAllProperties()
-    {
-        ImmutableMap.Builder builder = ImmutableMap.builder();
-        return builder.put(CACHE_ENABLED, String.valueOf(isCacheEnabled()))
-                .put(CACHE_MAX_CACHE_SIZE, String.valueOf(getDataSizeInLong(getMaxCacheSize().convertTo(DataSize.Unit.MEGABYTE))))
-                .put(CONNECTOR_NAME, getConnectorName())
-                .build();
-    }
-
-    public boolean isCacheEnabled()
-    {
-        return cacheEnabled;
-    }
-
-    @Config(CACHE_ENABLED)
-    public NativeExecutionConnectorConfig setCacheEnabled(boolean cacheEnabled)
-    {
-        this.cacheEnabled = cacheEnabled;
-        return this;
-    }
-
-    public DataSize getMaxCacheSize()
-    {
-        return maxCacheSize;
-    }
-
-    @Config(CACHE_MAX_CACHE_SIZE)
-    public NativeExecutionConnectorConfig setMaxCacheSize(DataSize maxCacheSize)
-    {
-        this.maxCacheSize = maxCacheSize;
-        return this;
-    }
-
-    public String getConnectorName()
-    {
-        return connectorName;
-    }
-
-    @Config(CONNECTOR_NAME)
-    public NativeExecutionConnectorConfig setConnectorName(String connectorName)
-    {
-        this.connectorName = connectorName;
-        return this;
-    }
-
-    private Long getDataSizeInLong(DataSize size)
-    {
-        return Double.valueOf(size.getValue()).longValue();
-    }
-}
diff --git a/presto-spark-base/src/main/java/com/facebook/presto/spark/execution/property/NativeExecutionSystemConfig.java b/presto-spark-base/src/main/java/com/facebook/presto/spark/execution/property/NativeExecutionSystemConfig.java
index f05eb2fa915c4..a0c757ab58595 100644
--- a/presto-spark-base/src/main/java/com/facebook/presto/spark/execution/property/NativeExecutionSystemConfig.java
+++ b/presto-spark-base/src/main/java/com/facebook/presto/spark/execution/property/NativeExecutionSystemConfig.java
@@ -13,740 +13,209 @@
  */
 package com.facebook.presto.spark.execution.property;
 
-import com.facebook.airlift.configuration.Config;
-import com.facebook.airlift.units.DataSize;
-import com.facebook.airlift.units.Duration;
 import com.google.common.collect.ImmutableMap;
+import org.apache.spark.SparkEnv$;
+import org.apache.spark.SparkFiles;
 
+import javax.inject.Inject;
+import javax.inject.Named;
+
+import java.io.File;
+import java.util.HashMap;
 import java.util.Map;
-import java.util.concurrent.TimeUnit;
 
+import static java.lang.String.format;
 import static java.util.Objects.requireNonNull;
 
 /**
- * This config class corresponds to config.properties for native execution process. Properties inside will be used in Configs::SystemConfig in Configs.h/cpp
+ * This config class corresponds to config.properties for native execution process. Properties
+ * inside will be used in Configs::SystemConfig in Configs.h/cpp
  */
 public class NativeExecutionSystemConfig
 {
-    private static final String CONCURRENT_LIFESPANS_PER_TASK = "concurrent-lifespans-per-task";
-    private static final String ENABLE_SERIALIZED_PAGE_CHECKSUM = "enable-serialized-page-checksum";
-    private static final String ENABLE_VELOX_EXPRESSION_LOGGING = "enable_velox_expression_logging";
-    private static final String ENABLE_VELOX_TASK_LOGGING = "enable_velox_task_logging";
-    // Port on which presto-native http server should run
-    private static final String HTTP_SERVER_HTTP_PORT = "http-server.http.port";
-    private static final String HTTP_SERVER_REUSE_PORT = "http-server.reuse-port";
-    private static final String HTTP_SERVER_BIND_TO_NODE_INTERNAL_ADDRESS_ONLY_ENABLED = "http-server.bind-to-node-internal-address-only-enabled";
-    // Number of I/O thread to use for serving http request on presto-native (proxygen server)
-    // this excludes worker thread used by velox
-    private static final String HTTP_SERVER_HTTPS_PORT = "http-server.https.port";
-    private static final String HTTP_SERVER_HTTPS_ENABLED = "http-server.https.enabled";
-
-    // This config control what cipher suites are supported by Native workers for server and client.
-    // Note Java and folly::SSLContext use different names to refer to the same cipher.
-    // (guess for different name, Java specific authentication,key exchange and cipher together and folly just cipher).
-    // For e.g. TLS_RSA_WITH_AES_256_GCM_SHA384 in Java and AES256-GCM-SHA384 in folly::SSLContext.
-    // The ciphers need to enable worker to worker, worker to coordinator and coordinator to worker communication.
-    // Have at least one cipher suite that is shared for the above 3, otherwise weird failures will result.
-    private static final String HTTPS_CIPHERS = "https-supported-ciphers";
-
-    // Note: Java packages cert and key in combined JKS file. But CPP requires them separately.
-    // The HTTPS provides integrity and not security(authentication/authorization).
-    // But the HTTPS will protect against data corruption by bad router and man in middle attacks.
-
-    // The cert path for the https server
-    private static final String HTTPS_CERT_PATH = "https-cert-path";
-    // The key path for the https server
-    private static final String HTTPS_KEY_PATH = "https-key-path";
-
-    // TODO: others use "-" separator and this property use _ separator. Fix them.
-    private static final String HTTP_SERVER_NUM_IO_THREADS_HW_MULTIPLIER = "http-server.num-io-threads-hw-multiplier";
-    private static final String EXCHANGE_HTTP_CLIENT_NUM_IO_THREADS_HW_MULTIPLIER = "exchange.http-client.num-io-threads-hw-multiplier";
-    private static final String ASYNC_DATA_CACHE_ENABLED = "async-data-cache-enabled";
-    private static final String ASYNC_CACHE_SSD_GB = "async-cache-ssd-gb";
-    private static final String CONNECTOR_NUM_IO_THREADS_HW_MULTIPLIER = "connector.num-io-threads-hw-multiplier";
-    private static final String PRESTO_VERSION = "presto.version";
-    private static final String SHUTDOWN_ONSET_SEC = "shutdown-onset-sec";
-    // Memory related configurations.
-    private static final String SYSTEM_MEMORY_GB = "system-memory-gb";
-    private static final String QUERY_MEMORY_GB = "query.max-memory-per-node";
-    private static final String USE_MMAP_ALLOCATOR = "use-mmap-allocator";
-    // Memory arbitration related configurations.
-    // Set the memory arbitrator kind. If it is empty, then there is no memory
-    // arbitration, when a query runs out of its capacity, the query will fail.
-    // If it set to "SHARED" (default), the shared memory arbitrator will be
-    // used to conduct arbitration and try to trigger disk spilling to reclaim
-    // memory so the query can run through completion.
-    private static final String MEMORY_ARBITRATOR_KIND = "memory-arbitrator-kind";
-    // Set memory arbitrator capacity to the same as per-query memory capacity
-    // as there is only one query running at Presto-on-Spark at a time.
-    private static final String MEMORY_ARBITRATOR_CAPACITY_GB = "query-memory-gb";
-    // Set memory arbitrator reserved capacity. Since there is only one query
-    // running at Presto-on-Spark at a time, then we shall set this to zero.
-    private static final String SHARED_ARBITRATOR_RESERVED_CAPACITY = "shared-arbitrator.reserved-capacity";
-    // Set the initial memory capacity when we create a query memory pool. For
-    // Presto-on-Spark, we set it to 'query-memory-gb' to allocate all the
-    // memory arbitrator capacity to the query memory pool on its creation as
-    // there is only one query running at a time.
-    private static final String SHARED_ARBITRATOR_MEMORY_POOL_INITIAL_CAPACITY = "shared-arbitrator.memory-pool-initial-capacity";
-    // Set the reserved memory capacity when we create a query memory pool. For
-    // Presto-on-Spark, we set this to zero as there is only one query running
-    // at a time.
-    private static final String SHARED_ARBITRATOR_MEMORY_POOL_RESERVED_CAPACITY = "shared-arbitrator.memory-pool-reserved-capacity";
-    private static final String SHARED_ARBITRATOR_MAX_MEMORY_ARBITRATION_TIME = "shared-arbitrator.max-memory-arbitration-time";
-    // Spilling related configs.
-    private static final String SPILLER_SPILL_PATH = "experimental.spiller-spill-path";
-    private static final String TASK_MAX_DRIVERS_PER_TASK = "task.max-drivers-per-task";
-    // Tasks are considered old, when they are in not-running state and it ended more than
-    // OLD_TASK_CLEANUP_MS ago or last heartbeat was more than OLD_TASK_CLEANUP_MS ago.
-    // For Presto-On-Spark, this is not relevant as it runs tasks serially, and spark's speculative
-    // execution takes care of zombie tasks.
-    private static final String ENABLE_OLD_TASK_CLEANUP = "enable-old-task-cleanup";
-    // Name of exchange client to use
-    private static final String SHUFFLE_NAME = "shuffle.name";
-    // Feature flag for access log on presto-native http server
-    private static final String HTTP_SERVER_ACCESS_LOGS = "http-server.enable-access-log";
-    // Terminates the native process and generates a core file on an allocation failure
-    private static final String CORE_ON_ALLOCATION_FAILURE_ENABLED = "core-on-allocation-failure-enabled";
-    // Spill related properties
-    private static final String SPILL_ENABLED = "spill-enabled";
-    private static final String AGGREGATION_SPILL_ENABLED = "aggregation-spill-enabled";
-    private static final String JOIN_SPILL_ENABLED = "join-spill-enabled";
-    private static final String ORDER_BY_SPILL_ENABLED = "order-by-spill-enabled";
-    private static final String MAX_SPILL_BYTES = "max-spill-bytes";
-
-    private boolean enableSerializedPageChecksum = true;
-    private boolean enableVeloxExpressionLogging;
-    private boolean enableVeloxTaskLogging = true;
-    private boolean httpServerReusePort = true;
-    private boolean httpServerBindToNodeInternalAddressOnlyEnabled = true;
-    private int httpServerPort = 7777;
-    private double httpServerNumIoThreadsHwMultiplier = 1.0;
-    private int httpsServerPort = 7778;
-    private boolean enableHttpsCommunication;
-    private String httpsCiphers = "AES128-SHA,AES128-SHA256,AES256-GCM-SHA384";
-    private String httpsCertPath = "";
-    private String httpsKeyPath = "";
-    private double exchangeHttpClientNumIoThreadsHwMultiplier = 1.0;
-    private boolean asyncDataCacheEnabled; // false
-    private int asyncCacheSsdGb; // 0
-    private double connectorNumIoThreadsHwMultiplier; // 0.0
-    private int shutdownOnsetSec = 10;
-    private int systemMemoryGb = 10;
-    // Reserve 2GB from system memory for system operations such as disk
-    // spilling and cache prefetch.
-    private DataSize queryMemoryGb = new DataSize(8, DataSize.Unit.GIGABYTE);
-
-    private boolean useMmapAllocator = true;
-    private String memoryArbitratorKind = "SHARED";
-    private int memoryArbitratorCapacityGb = 8;
-    private DataSize sharedArbitratorReservedCapacity = new DataSize(0, DataSize.Unit.GIGABYTE);
-    private DataSize sharedArbitratorMemoryPoolInitialCapacity = new DataSize(4, DataSize.Unit.GIGABYTE);
-    private DataSize sharedArbitratorMemoryPoolReservedCapacity = new DataSize(32, DataSize.Unit.MEGABYTE);
-    private Duration sharedArbitratorMaxMemoryArbitrationTime = new Duration(5, TimeUnit.MINUTES);
-    private String spillerSpillPath = "";
-    private int concurrentLifespansPerTask = 5;
-    private int maxDriversPerTask = 15;
-    private boolean enableOldTaskCleanUp; // false;
-    private String prestoVersion = "dummy.presto.version";
-    private String shuffleName = "local";
-    private boolean enableHttpServerAccessLog = true;
-    private boolean coreOnAllocationFailureEnabled;
-    private boolean spillEnabled = true;
-    private boolean aggregationSpillEnabled = true;
-    private boolean joinSpillEnabled = true;
-    private boolean orderBySpillEnabled = true;
-    private Long maxSpillBytes = 600L << 30;
-
-    // TODO: Deprecate following configs
-    private static final String REGISTER_TEST_FUNCTIONS = "register-test-functions";
-    private static final String MEMORY_POOL_INIT_CAPACITY = "memory-pool-init-capacity";
-    private static final String MEMORY_POOL_RESERVED_CAPACITY = "memory-pool-reserved-capacity";
-    private static final String MEMORY_POOL_TRANSFER_CAPACITY = "memory-pool-transfer-capacity";
-    private static final String MEMORY_RECLAIM_WAIT_MS = "memory-reclaim-wait-ms";
-    private boolean registerTestFunctions;
-    private long memoryPoolInitCapacity;
-    private long memoryPoolReservedCapacity;
-    private long memoryPoolTransferCapacity;
-    private long memoryReclaimWaitMs;
-
-    public Map getAllProperties()
-    {
-        ImmutableMap.Builder builder = ImmutableMap.builder();
-        return builder.put(CONCURRENT_LIFESPANS_PER_TASK, String.valueOf(getConcurrentLifespansPerTask()))
-                .put(ENABLE_SERIALIZED_PAGE_CHECKSUM, String.valueOf(isEnableSerializedPageChecksum()))
-                .put(ENABLE_VELOX_EXPRESSION_LOGGING, String.valueOf(isEnableVeloxExpressionLogging()))
-                .put(ENABLE_VELOX_TASK_LOGGING, String.valueOf(isEnableVeloxTaskLogging()))
-                .put(HTTP_SERVER_HTTP_PORT, String.valueOf(getHttpServerPort()))
-                .put(HTTP_SERVER_REUSE_PORT, String.valueOf(isHttpServerReusePort()))
-                .put(HTTP_SERVER_BIND_TO_NODE_INTERNAL_ADDRESS_ONLY_ENABLED, String.valueOf(isHttpServerBindToNodeInternalAddressOnlyEnabled()))
-                .put(HTTP_SERVER_HTTPS_PORT, String.valueOf(getHttpsServerPort()))
-                .put(HTTP_SERVER_HTTPS_ENABLED, String.valueOf(isEnableHttpsCommunication()))
-                .put(HTTPS_CIPHERS, String.valueOf(getHttpsCiphers()))
-                .put(HTTPS_CERT_PATH, String.valueOf(getHttpsCertPath()))
-                .put(HTTPS_KEY_PATH, String.valueOf(getHttpsKeyPath()))
-                .put(HTTP_SERVER_NUM_IO_THREADS_HW_MULTIPLIER, String.valueOf(getHttpServerNumIoThreadsHwMultiplier()))
-                .put(EXCHANGE_HTTP_CLIENT_NUM_IO_THREADS_HW_MULTIPLIER, String.valueOf(getExchangeHttpClientNumIoThreadsHwMultiplier()))
-                .put(ASYNC_DATA_CACHE_ENABLED, String.valueOf(getAsyncDataCacheEnabled()))
-                .put(ASYNC_CACHE_SSD_GB, String.valueOf(getAsyncCacheSsdGb()))
-                .put(CONNECTOR_NUM_IO_THREADS_HW_MULTIPLIER, String.valueOf(getConnectorNumIoThreadsHwMultiplier()))
-                .put(PRESTO_VERSION, getPrestoVersion())
-                .put(SHUTDOWN_ONSET_SEC, String.valueOf(getShutdownOnsetSec()))
-                .put(SYSTEM_MEMORY_GB, String.valueOf(getSystemMemoryGb()))
-                .put(QUERY_MEMORY_GB, String.valueOf(getQueryMemoryGb()))
-                .put(USE_MMAP_ALLOCATOR, String.valueOf(getUseMmapAllocator()))
-                .put(MEMORY_ARBITRATOR_KIND, String.valueOf(getMemoryArbitratorKind()))
-                .put(MEMORY_ARBITRATOR_CAPACITY_GB, String.valueOf(getMemoryArbitratorCapacityGb()))
-                .put(SHARED_ARBITRATOR_RESERVED_CAPACITY, String.valueOf(getSharedArbitratorReservedCapacity()))
-                .put(SHARED_ARBITRATOR_MEMORY_POOL_INITIAL_CAPACITY, String.valueOf(getSharedArbitratorMemoryPoolInitialCapacity()))
-                .put(SHARED_ARBITRATOR_MEMORY_POOL_RESERVED_CAPACITY, String.valueOf(getSharedArbitratorMemoryPoolReservedCapacity()))
-                .put(SHARED_ARBITRATOR_MAX_MEMORY_ARBITRATION_TIME, String.valueOf(getSharedArbitratorMaxMemoryArbitrationTime()))
-                .put(SPILLER_SPILL_PATH, String.valueOf(getSpillerSpillPath()))
-                .put(TASK_MAX_DRIVERS_PER_TASK, String.valueOf(getMaxDriversPerTask()))
-                .put(ENABLE_OLD_TASK_CLEANUP, String.valueOf(getOldTaskCleanupMs()))
-                .put(SHUFFLE_NAME, getShuffleName())
-                .put(HTTP_SERVER_ACCESS_LOGS, String.valueOf(isEnableHttpServerAccessLog()))
-                .put(CORE_ON_ALLOCATION_FAILURE_ENABLED, String.valueOf(isCoreOnAllocationFailureEnabled()))
-                .put(SPILL_ENABLED, String.valueOf(getSpillEnabled()))
-                .put(AGGREGATION_SPILL_ENABLED, String.valueOf(getAggregationSpillEnabled()))
-                .put(JOIN_SPILL_ENABLED, String.valueOf(getJoinSpillEnabled()))
-                .put(ORDER_BY_SPILL_ENABLED, String.valueOf(getOrderBySpillEnabled()))
-                .put(MAX_SPILL_BYTES, String.valueOf(getMaxSpillBytes()))
-                .put(REGISTER_TEST_FUNCTIONS, String.valueOf(registerTestFunctions))
-                .put(MEMORY_POOL_INIT_CAPACITY, String.valueOf(memoryPoolInitCapacity))
-                .put(MEMORY_POOL_RESERVED_CAPACITY, String.valueOf(memoryPoolReservedCapacity))
-                .put(MEMORY_POOL_TRANSFER_CAPACITY, String.valueOf(memoryPoolTransferCapacity))
-                .put(MEMORY_RECLAIM_WAIT_MS, String.valueOf(memoryReclaimWaitMs))
+    public static final String NATIVE_EXECUTION_SYSTEM_CONFIG = "native-execution-system-config";
+
+    public static final String CONCURRENT_LIFESPANS_PER_TASK = "concurrent-lifespans-per-task";
+    public static final String ENABLE_SERIALIZED_PAGE_CHECKSUM = "enable-serialized-page-checksum";
+    public static final String ENABLE_VELOX_EXPRESSION_LOGGING = "enable_velox_expression_logging";
+    public static final String ENABLE_VELOX_TASK_LOGGING = "enable_velox_task_logging";
+    public static final String HTTP_SERVER_HTTP_PORT = "http-server.http.port";
+    public static final String HTTP_SERVER_REUSE_PORT = "http-server.reuse-port";
+    public static final String HTTP_SERVER_BIND_TO_NODE_INTERNAL_ADDRESS_ONLY_ENABLED = "http-server.bind-to-node-internal-address-only-enabled";
+    public static final String HTTP_SERVER_HTTPS_PORT = "http-server.https.port";
+    public static final String HTTP_SERVER_HTTPS_ENABLED = "http-server.https.enabled";
+    public static final String HTTPS_SUPPORTED_CIPHERS = "https-supported-ciphers";
+    public static final String HTTPS_CERT_PATH = "https-cert-path";
+    public static final String HTTPS_KEY_PATH = "https-key-path";
+    public static final String HTTP_SERVER_NUM_IO_THREADS_HW_MULTIPLIER = "http-server.num-io-threads-hw-multiplier";
+    public static final String EXCHANGE_HTTP_CLIENT_NUM_IO_THREADS_HW_MULTIPLIER = "exchange.http-client.num-io-threads-hw-multiplier";
+    public static final String ASYNC_DATA_CACHE_ENABLED = "async-data-cache-enabled";
+    public static final String ASYNC_CACHE_SSD_GB = "async-cache-ssd-gb";
+    public static final String CONNECTOR_NUM_IO_THREADS_HW_MULTIPLIER = "connector.num-io-threads-hw-multiplier";
+    public static final String PRESTO_VERSION = "presto.version";
+    public static final String SHUTDOWN_ONSET_SEC = "shutdown-onset-sec";
+    public static final String SYSTEM_MEMORY_GB = "system-memory-gb";
+    public static final String QUERY_MEMORY_GB = "query-memory-gb";
+    public static final String QUERY_MAX_MEMORY_PER_NODE = "query.max-memory-per-node";
+    public static final String USE_MMAP_ALLOCATOR = "use-mmap-allocator";
+    public static final String MEMORY_ARBITRATOR_KIND = "memory-arbitrator-kind";
+    public static final String SHARED_ARBITRATOR_RESERVED_CAPACITY = "shared-arbitrator.reserved-capacity";
+    public static final String SHARED_ARBITRATOR_MEMORY_POOL_INITIAL_CAPACITY = "shared-arbitrator.memory-pool-initial-capacity";
+    public static final String SHARED_ARBITRATOR_MAX_MEMORY_ARBITRATION_TIME = "shared-arbitrator.max-memory-arbitration-time";
+    public static final String EXPERIMENTAL_SPILLER_SPILL_PATH = "experimental.spiller-spill-path";
+    public static final String TASK_MAX_DRIVERS_PER_TASK = "task.max-drivers-per-task";
+    public static final String ENABLE_OLD_TASK_CLEANUP = "enable-old-task-cleanup";
+    public static final String SHUFFLE_NAME = "shuffle.name";
+    public static final String HTTP_SERVER_ENABLE_ACCESS_LOG = "http-server.enable-access-log";
+    public static final String CORE_ON_ALLOCATION_FAILURE_ENABLED = "core-on-allocation-failure-enabled";
+    public static final String SPILL_ENABLED = "spill-enabled";
+    public static final String AGGREGATION_SPILL_ENABLED = "aggregation-spill-enabled";
+    public static final String JOIN_SPILL_ENABLED = "join-spill-enabled";
+    public static final String ORDER_BY_SPILL_ENABLED = "order-by-spill-enabled";
+    public static final String MAX_SPILL_BYTES = "max-spill-bytes";
+    public static final String REMOTE_FUNCTION_SERVER_THRIFT_UDS_PATH = "remote-function-server.thrift.uds-path";
+    public static final String REMOTE_FUNCTION_SERVER_SIGNATURE_FILES_DIRECTORY_PATH = "remote-function-server.signature.files.directory.path";
+    public static final String REMOTE_FUNCTION_SERVER_SERDE = "remote-function-server.serde";
+    public static final String REMOTE_FUNCTION_SERVER_CATALOG_NAME = "remote-function-server.catalog-name";
+
+    private final String remoteFunctionServerSignatureFilesDirectoryPathDefault = "./functions/spark/";
+    private final String remoteFunctionServerSerdeDefault = "presto_page";
+    private final String remoteFunctionServerCatalogNameDefault = "";
+    private final String concurrentLifespansPerTaskDefault = "5";
+    private final String enableSerializedPageChecksumDefault = "true";
+    private final String enableVeloxExpressionLoggingDefault = "false";
+    private final String enableVeloxTaskLoggingDefault = "true";
+    private final String httpServerHttpPortDefault = "7777";
+    private final String httpServerReusePortDefault = "true";
+    private final String httpServerBindToNodeInternalAddressOnlyEnabledDefault = "true";
+    private final String httpServerHttpsPortDefault = "7778";
+    private final String httpServerHttpsEnabledDefault = "false";
+    private final String httpsSupportedCiphersDefault = "AES128-SHA,AES128-SHA256,AES256-GCM-SHA384";
+    private final String httpsCertPathDefault = "";
+    private final String httpsKeyPathDefault = "";
+    private final String httpServerNumIoThreadsHwMultiplierDefault = "1.0";
+    private final String exchangeHttpClientNumIoThreadsHwMultiplierDefault = "1.0";
+    private final String asyncDataCacheEnabledDefault = "false";
+    private final String asyncCacheSsdGbDefault = "0";
+    private final String connectorNumIoThreadsHwMultiplierDefault = "0";
+    private final String prestoVersionDefault = "dummy.presto.version";
+    private final String shutdownOnsetSecDefault = "10";
+    private final String systemMemoryGbDefault = "10";
+    private final String queryMemoryGbDefault = "8";
+    private final String queryMaxMemoryPerNodeDefault = "8GB";
+    private final String useMmapAllocatorDefault = "true";
+    private final String memoryArbitratorKindDefault = "SHARED";
+    private final String sharedArbitratorReservedCapacityDefault = "0GB";
+    private final String sharedArbitratorMemoryPoolInitialCapacityDefault = "4GB";
+    private final String sharedArbitratorMaxMemoryArbitrationTimeDefault = "5m";
+    private final String experimentalSpillerSpillPathDefault = "";
+    private final String taskMaxDriversPerTaskDefault = "15";
+    private final String enableOldTaskCleanupDefault = "false";
+    private final String shuffleNameDefault = "local";
+    private final String httpServerEnableAccessLogDefault = "true";
+    private final String coreOnAllocationFailureEnabledDefault = "false";
+    private final String spillEnabledDefault = "true";
+    private final String aggregationSpillEnabledDefault = "true";
+    private final String joinSpillEnabledDefault = "true";
+    private final String orderBySpillEnabledDefault = "true";
+    private final String maxSpillBytesDefault = String.valueOf(600L << 30);
+
+    private final Map systemConfigs;
+    private final Map defaultSystemConfigs;
+
+    @Inject
+    public NativeExecutionSystemConfig(
+            @Named(NATIVE_EXECUTION_SYSTEM_CONFIG) Map systemConfigs)
+    {
+        this.systemConfigs = new HashMap<>(
+                requireNonNull(systemConfigs, "systemConfigs is null"));
+
+        ImmutableMap.Builder defaultSystemConfigsBuilder = ImmutableMap.builder();
+
+        String remoteFunctionServerThriftUdsPath = System.getProperty(REMOTE_FUNCTION_SERVER_THRIFT_UDS_PATH);
+        if (remoteFunctionServerThriftUdsPath != null) {
+            defaultSystemConfigsBuilder.put(REMOTE_FUNCTION_SERVER_THRIFT_UDS_PATH, remoteFunctionServerThriftUdsPath);
+            defaultSystemConfigsBuilder.put(REMOTE_FUNCTION_SERVER_SIGNATURE_FILES_DIRECTORY_PATH, getAbsolutePath(remoteFunctionServerSignatureFilesDirectoryPathDefault));
+            defaultSystemConfigsBuilder.put(REMOTE_FUNCTION_SERVER_SERDE, remoteFunctionServerSerdeDefault);
+            defaultSystemConfigsBuilder.put(REMOTE_FUNCTION_SERVER_CATALOG_NAME, remoteFunctionServerCatalogNameDefault);
+        }
+
+        defaultSystemConfigs = defaultSystemConfigsBuilder
+                .put(CONCURRENT_LIFESPANS_PER_TASK, concurrentLifespansPerTaskDefault)
+                .put(ENABLE_SERIALIZED_PAGE_CHECKSUM, enableSerializedPageChecksumDefault)
+                .put(ENABLE_VELOX_EXPRESSION_LOGGING, enableVeloxExpressionLoggingDefault)
+                .put(ENABLE_VELOX_TASK_LOGGING, enableVeloxTaskLoggingDefault)
+                .put(HTTP_SERVER_HTTP_PORT, httpServerHttpPortDefault)
+                .put(HTTP_SERVER_REUSE_PORT, httpServerReusePortDefault)
+                .put(HTTP_SERVER_BIND_TO_NODE_INTERNAL_ADDRESS_ONLY_ENABLED,
+                    httpServerBindToNodeInternalAddressOnlyEnabledDefault)
+                .put(HTTP_SERVER_HTTPS_PORT, httpServerHttpsPortDefault)
+                .put(HTTP_SERVER_HTTPS_ENABLED, httpServerHttpsEnabledDefault)
+                .put(HTTPS_SUPPORTED_CIPHERS, httpsSupportedCiphersDefault)
+                .put(HTTPS_CERT_PATH, httpsCertPathDefault)
+                .put(HTTPS_KEY_PATH, httpsKeyPathDefault)
+                .put(HTTP_SERVER_NUM_IO_THREADS_HW_MULTIPLIER,
+                    httpServerNumIoThreadsHwMultiplierDefault)
+                .put(EXCHANGE_HTTP_CLIENT_NUM_IO_THREADS_HW_MULTIPLIER,
+                    exchangeHttpClientNumIoThreadsHwMultiplierDefault)
+                .put(ASYNC_DATA_CACHE_ENABLED, asyncDataCacheEnabledDefault)
+                .put(ASYNC_CACHE_SSD_GB, asyncCacheSsdGbDefault)
+                .put(CONNECTOR_NUM_IO_THREADS_HW_MULTIPLIER, connectorNumIoThreadsHwMultiplierDefault)
+                .put(PRESTO_VERSION, prestoVersionDefault)
+                .put(SHUTDOWN_ONSET_SEC, shutdownOnsetSecDefault)
+                .put(SYSTEM_MEMORY_GB, systemMemoryGbDefault)
+                .put(QUERY_MEMORY_GB, queryMemoryGbDefault)
+                .put(QUERY_MAX_MEMORY_PER_NODE, queryMaxMemoryPerNodeDefault)
+                .put(USE_MMAP_ALLOCATOR, useMmapAllocatorDefault)
+                .put(MEMORY_ARBITRATOR_KIND, memoryArbitratorKindDefault)
+                .put(SHARED_ARBITRATOR_RESERVED_CAPACITY, sharedArbitratorReservedCapacityDefault)
+                .put(SHARED_ARBITRATOR_MEMORY_POOL_INITIAL_CAPACITY,
+                    sharedArbitratorMemoryPoolInitialCapacityDefault)
+                .put(SHARED_ARBITRATOR_MAX_MEMORY_ARBITRATION_TIME,
+                    sharedArbitratorMaxMemoryArbitrationTimeDefault)
+                .put(EXPERIMENTAL_SPILLER_SPILL_PATH, experimentalSpillerSpillPathDefault)
+                .put(TASK_MAX_DRIVERS_PER_TASK, taskMaxDriversPerTaskDefault)
+                .put(ENABLE_OLD_TASK_CLEANUP, enableOldTaskCleanupDefault)
+                .put(SHUFFLE_NAME, shuffleNameDefault)
+                .put(HTTP_SERVER_ENABLE_ACCESS_LOG, httpServerEnableAccessLogDefault)
+                .put(CORE_ON_ALLOCATION_FAILURE_ENABLED, coreOnAllocationFailureEnabledDefault)
+                .put(SPILL_ENABLED, spillEnabledDefault)
+                .put(AGGREGATION_SPILL_ENABLED, aggregationSpillEnabledDefault)
+                .put(JOIN_SPILL_ENABLED, joinSpillEnabledDefault)
+                .put(ORDER_BY_SPILL_ENABLED, orderBySpillEnabledDefault)
+                .put(MAX_SPILL_BYTES, maxSpillBytesDefault)
                 .build();
     }
 
-    @Config(SHUFFLE_NAME)
-    public NativeExecutionSystemConfig setShuffleName(String shuffleName)
-    {
-        this.shuffleName = requireNonNull(shuffleName);
-        return this;
-    }
-
-    public String getShuffleName()
-    {
-        return shuffleName;
-    }
-
-    @Config(ENABLE_SERIALIZED_PAGE_CHECKSUM)
-    public NativeExecutionSystemConfig setEnableSerializedPageChecksum(boolean enableSerializedPageChecksum)
-    {
-        this.enableSerializedPageChecksum = enableSerializedPageChecksum;
-        return this;
-    }
-
-    public boolean isEnableSerializedPageChecksum()
-    {
-        return enableSerializedPageChecksum;
-    }
-
-    @Config(ENABLE_VELOX_EXPRESSION_LOGGING)
-    public NativeExecutionSystemConfig setEnableVeloxExpressionLogging(boolean enableVeloxExpressionLogging)
-    {
-        this.enableVeloxExpressionLogging = enableVeloxExpressionLogging;
-        return this;
-    }
-
-    public boolean isEnableVeloxExpressionLogging()
-    {
-        return enableVeloxExpressionLogging;
-    }
-
-    @Config(ENABLE_VELOX_TASK_LOGGING)
-    public NativeExecutionSystemConfig setEnableVeloxTaskLogging(boolean enableVeloxTaskLogging)
-    {
-        this.enableVeloxTaskLogging = enableVeloxTaskLogging;
-        return this;
-    }
-
-    public boolean isEnableVeloxTaskLogging()
-    {
-        return enableVeloxTaskLogging;
-    }
-
-    @Config(HTTP_SERVER_HTTP_PORT)
-    public NativeExecutionSystemConfig setHttpServerPort(int httpServerPort)
-    {
-        this.httpServerPort = httpServerPort;
-        return this;
-    }
-
-    public int getHttpServerPort()
-    {
-        return httpServerPort;
-    }
-
-    @Config(HTTP_SERVER_REUSE_PORT)
-    public NativeExecutionSystemConfig setHttpServerReusePort(boolean httpServerReusePort)
-    {
-        this.httpServerReusePort = httpServerReusePort;
-        return this;
-    }
-
-    public boolean isHttpServerReusePort()
-    {
-        return httpServerReusePort;
-    }
-
-    public boolean isHttpServerBindToNodeInternalAddressOnlyEnabled()
-    {
-        return httpServerBindToNodeInternalAddressOnlyEnabled;
-    }
-
-    @Config(HTTP_SERVER_BIND_TO_NODE_INTERNAL_ADDRESS_ONLY_ENABLED)
-    public NativeExecutionSystemConfig setHttpServerBindToNodeInternalAddressOnlyEnabled(boolean httpServerBindToNodeInternalAddressOnlyEnabled)
-    {
-        this.httpServerBindToNodeInternalAddressOnlyEnabled = httpServerBindToNodeInternalAddressOnlyEnabled;
-        return this;
-    }
-
-    @Config(HTTP_SERVER_NUM_IO_THREADS_HW_MULTIPLIER)
-    public NativeExecutionSystemConfig setHttpServerNumIoThreadsHwMultiplier(double httpServerNumIoThreadsHwMultiplier)
-    {
-        this.httpServerNumIoThreadsHwMultiplier = httpServerNumIoThreadsHwMultiplier;
-        return this;
-    }
-
-    public double getHttpServerNumIoThreadsHwMultiplier()
-    {
-        return httpServerNumIoThreadsHwMultiplier;
-    }
-
-    public int getHttpsServerPort()
-    {
-        return httpsServerPort;
-    }
-
-    @Config(HTTP_SERVER_HTTPS_PORT)
-    public NativeExecutionSystemConfig setHttpsServerPort(int httpsServerPort)
-    {
-        this.httpsServerPort = httpsServerPort;
-        return this;
-    }
-
-    public boolean isEnableHttpsCommunication()
-    {
-        return enableHttpsCommunication;
-    }
-
-    @Config(HTTP_SERVER_HTTPS_ENABLED)
-    public NativeExecutionSystemConfig setEnableHttpsCommunication(boolean enableHttpsCommunication)
-    {
-        this.enableHttpsCommunication = enableHttpsCommunication;
-        return this;
-    }
-
-    public String getHttpsCiphers()
-    {
-        return httpsCiphers;
-    }
-
-    @Config(HTTPS_CIPHERS)
-    public NativeExecutionSystemConfig setHttpsCiphers(String httpsCiphers)
-    {
-        this.httpsCiphers = httpsCiphers;
-        return this;
-    }
-
-    public String getHttpsCertPath()
-    {
-        return httpsCertPath;
-    }
-
-    @Config(HTTPS_CERT_PATH)
-    public NativeExecutionSystemConfig setHttpsCertPath(String httpsCertPath)
-    {
-        this.httpsCertPath = httpsCertPath;
-        return this;
-    }
-
-    public String getHttpsKeyPath()
-    {
-        return httpsKeyPath;
-    }
-
-    @Config(HTTPS_KEY_PATH)
-    public NativeExecutionSystemConfig setHttpsKeyPath(String httpsKeyPath)
-    {
-        this.httpsKeyPath = httpsKeyPath;
-        return this;
-    }
-
-    @Config(EXCHANGE_HTTP_CLIENT_NUM_IO_THREADS_HW_MULTIPLIER)
-    public NativeExecutionSystemConfig setExchangeHttpClientNumIoThreadsHwMultiplier(double exchangeHttpClientNumIoThreadsHwMultiplier)
-    {
-        this.exchangeHttpClientNumIoThreadsHwMultiplier = exchangeHttpClientNumIoThreadsHwMultiplier;
-        return this;
-    }
-
-    public double getExchangeHttpClientNumIoThreadsHwMultiplier()
-    {
-        return exchangeHttpClientNumIoThreadsHwMultiplier;
-    }
-
-    @Config(ASYNC_DATA_CACHE_ENABLED)
-    public NativeExecutionSystemConfig setAsyncDataCacheEnabled(boolean asyncDataCacheEnabled)
-    {
-        this.asyncDataCacheEnabled = asyncDataCacheEnabled;
-        return this;
-    }
-
-    public boolean getAsyncDataCacheEnabled()
-    {
-        return asyncDataCacheEnabled;
-    }
-
-    @Config(ASYNC_CACHE_SSD_GB)
-    public NativeExecutionSystemConfig setAsyncCacheSsdGb(int asyncCacheSsdGb)
-    {
-        this.asyncCacheSsdGb = asyncCacheSsdGb;
-        return this;
-    }
-
-    public int getAsyncCacheSsdGb()
-    {
-        return asyncCacheSsdGb;
-    }
-
-    @Config(CONNECTOR_NUM_IO_THREADS_HW_MULTIPLIER)
-    public NativeExecutionSystemConfig setConnectorNumIoThreadsHwMultiplier(double connectorNumIoThreadsHwMultiplier)
-    {
-        this.connectorNumIoThreadsHwMultiplier = connectorNumIoThreadsHwMultiplier;
-        return this;
-    }
-
-    public double getConnectorNumIoThreadsHwMultiplier()
-    {
-        return connectorNumIoThreadsHwMultiplier;
-    }
-
-    @Config(SHUTDOWN_ONSET_SEC)
-    public NativeExecutionSystemConfig setShutdownOnsetSec(int shutdownOnsetSec)
-    {
-        this.shutdownOnsetSec = shutdownOnsetSec;
-        return this;
-    }
-
-    public int getShutdownOnsetSec()
-    {
-        return shutdownOnsetSec;
-    }
-
-    @Config(SYSTEM_MEMORY_GB)
-    public NativeExecutionSystemConfig setSystemMemoryGb(int systemMemoryGb)
-    {
-        this.systemMemoryGb = systemMemoryGb;
-        return this;
-    }
-
-    public int getSystemMemoryGb()
-    {
-        return systemMemoryGb;
-    }
-
-    @Config(QUERY_MEMORY_GB)
-    public NativeExecutionSystemConfig setQueryMemoryGb(DataSize queryMemoryGb)
-    {
-        this.queryMemoryGb = queryMemoryGb;
-        return this;
-    }
-
-    public DataSize getQueryMemoryGb()
-    {
-        return queryMemoryGb;
-    }
-
-    @Config(USE_MMAP_ALLOCATOR)
-    public NativeExecutionSystemConfig setUseMmapAllocator(boolean useMmapAllocator)
-    {
-        this.useMmapAllocator = useMmapAllocator;
-        return this;
-    }
-
-    public boolean getUseMmapAllocator()
-    {
-        return useMmapAllocator;
-    }
-
-    @Config(MEMORY_ARBITRATOR_KIND)
-    public NativeExecutionSystemConfig setMemoryArbitratorKind(String memoryArbitratorKind)
-    {
-        this.memoryArbitratorKind = memoryArbitratorKind;
-        return this;
-    }
-
-    public String getMemoryArbitratorKind()
-    {
-        return memoryArbitratorKind;
-    }
-
-    @Config(MEMORY_ARBITRATOR_CAPACITY_GB)
-    public NativeExecutionSystemConfig setMemoryArbitratorCapacityGb(int memoryArbitratorCapacityGb)
-    {
-        this.memoryArbitratorCapacityGb = memoryArbitratorCapacityGb;
-        return this;
-    }
-
-    public int getMemoryArbitratorCapacityGb()
-    {
-        return memoryArbitratorCapacityGb;
-    }
-
-    @Config(SHARED_ARBITRATOR_RESERVED_CAPACITY)
-    public NativeExecutionSystemConfig setSharedArbitratorReservedCapacity(DataSize sharedArbitratorReservedCapacity)
-    {
-        this.sharedArbitratorReservedCapacity = sharedArbitratorReservedCapacity;
-        return this;
-    }
-
-    public DataSize getSharedArbitratorReservedCapacity()
-    {
-        return sharedArbitratorReservedCapacity;
-    }
-
-    @Config(SHARED_ARBITRATOR_MEMORY_POOL_INITIAL_CAPACITY)
-    public NativeExecutionSystemConfig setSharedArbitratorMemoryPoolInitialCapacity(DataSize sharedArbitratorMemoryPoolInitialCapacity)
-    {
-        this.sharedArbitratorMemoryPoolInitialCapacity = sharedArbitratorMemoryPoolInitialCapacity;
-        return this;
-    }
-
-    public DataSize getSharedArbitratorMemoryPoolInitialCapacity()
-    {
-        return sharedArbitratorMemoryPoolInitialCapacity;
-    }
-
-    @Config(SHARED_ARBITRATOR_MEMORY_POOL_RESERVED_CAPACITY)
-    public NativeExecutionSystemConfig setSharedArbitratorMemoryPoolReservedCapacity(DataSize sharedArbitratorMemoryPoolReservedCapacity)
-    {
-        this.sharedArbitratorMemoryPoolReservedCapacity = sharedArbitratorMemoryPoolReservedCapacity;
-        return this;
-    }
-
-    public DataSize getSharedArbitratorMemoryPoolReservedCapacity()
-    {
-        return sharedArbitratorMemoryPoolReservedCapacity;
-    }
-
-    @Config(SHARED_ARBITRATOR_MAX_MEMORY_ARBITRATION_TIME)
-    public NativeExecutionSystemConfig setSharedArbitratorMaxMemoryArbitrationTime(Duration sharedArbitratorMaxMemoryArbitrationTime)
-    {
-        this.sharedArbitratorMaxMemoryArbitrationTime = sharedArbitratorMaxMemoryArbitrationTime;
-        return this;
-    }
-
-    public Duration getSharedArbitratorMaxMemoryArbitrationTime()
-    {
-        return sharedArbitratorMaxMemoryArbitrationTime;
-    }
-
-    @Config(SPILLER_SPILL_PATH)
-    public NativeExecutionSystemConfig setSpillerSpillPath(String spillerSpillPath)
-    {
-        this.spillerSpillPath = spillerSpillPath;
-        return this;
-    }
-
-    public String getSpillerSpillPath()
-    {
-        return spillerSpillPath;
-    }
-
-    @Config(CONCURRENT_LIFESPANS_PER_TASK)
-    public NativeExecutionSystemConfig setConcurrentLifespansPerTask(int concurrentLifespansPerTask)
-    {
-        this.concurrentLifespansPerTask = concurrentLifespansPerTask;
-        return this;
-    }
-
-    public int getConcurrentLifespansPerTask()
-    {
-        return concurrentLifespansPerTask;
-    }
-
-    @Config(TASK_MAX_DRIVERS_PER_TASK)
-    public NativeExecutionSystemConfig setMaxDriversPerTask(int maxDriversPerTask)
-    {
-        this.maxDriversPerTask = maxDriversPerTask;
-        return this;
-    }
-
-    public int getMaxDriversPerTask()
-    {
-        return maxDriversPerTask;
-    }
-
-    public boolean getOldTaskCleanupMs()
-    {
-        return enableOldTaskCleanUp;
-    }
-
-    @Config(ENABLE_OLD_TASK_CLEANUP)
-    public NativeExecutionSystemConfig setOldTaskCleanupMs(boolean enableOldTaskCleanUp)
-    {
-        this.enableOldTaskCleanUp = enableOldTaskCleanUp;
-        return this;
-    }
-
-    @Config(PRESTO_VERSION)
-    public NativeExecutionSystemConfig setPrestoVersion(String prestoVersion)
-    {
-        this.prestoVersion = prestoVersion;
-        return this;
-    }
-
-    public String getPrestoVersion()
-    {
-        return prestoVersion;
-    }
-
-    @Config(HTTP_SERVER_ACCESS_LOGS)
-    public NativeExecutionSystemConfig setEnableHttpServerAccessLog(boolean enableHttpServerAccessLog)
-    {
-        this.enableHttpServerAccessLog = enableHttpServerAccessLog;
-        return this;
-    }
-
-    public boolean isEnableHttpServerAccessLog()
-    {
-        return enableHttpServerAccessLog;
-    }
-
-    public boolean isCoreOnAllocationFailureEnabled()
-    {
-        return coreOnAllocationFailureEnabled;
-    }
-
-    @Config(CORE_ON_ALLOCATION_FAILURE_ENABLED)
-    public NativeExecutionSystemConfig setCoreOnAllocationFailureEnabled(boolean coreOnAllocationFailureEnabled)
-    {
-        this.coreOnAllocationFailureEnabled = coreOnAllocationFailureEnabled;
-        return this;
-    }
-
-    public boolean getSpillEnabled()
-    {
-        return spillEnabled;
-    }
-
-    @Config(SPILL_ENABLED)
-    public NativeExecutionSystemConfig setSpillEnabled(boolean spillEnabled)
-    {
-        this.spillEnabled = spillEnabled;
-        return this;
-    }
-
-    public boolean getAggregationSpillEnabled()
-    {
-        return aggregationSpillEnabled;
-    }
-
-    @Config(AGGREGATION_SPILL_ENABLED)
-    public NativeExecutionSystemConfig setAggregationSpillEnabled(boolean aggregationSpillEnabled)
-    {
-        this.aggregationSpillEnabled = aggregationSpillEnabled;
-        return this;
-    }
-
-    public boolean getJoinSpillEnabled()
-    {
-        return joinSpillEnabled;
-    }
-
-    @Config(JOIN_SPILL_ENABLED)
-    public NativeExecutionSystemConfig setJoinSpillEnabled(boolean joinSpillEnabled)
-    {
-        this.joinSpillEnabled = joinSpillEnabled;
-        return this;
-    }
-
-    public boolean getOrderBySpillEnabled()
-    {
-        return orderBySpillEnabled;
-    }
-
-    @Config(ORDER_BY_SPILL_ENABLED)
-    public NativeExecutionSystemConfig setOrderBySpillEnabled(boolean orderBySpillEnabled)
-    {
-        this.orderBySpillEnabled = orderBySpillEnabled;
-        return this;
-    }
-
-    public Long getMaxSpillBytes()
-    {
-        return maxSpillBytes;
-    }
-
-    @Config(MAX_SPILL_BYTES)
-    public NativeExecutionSystemConfig setMaxSpillBytes(Long maxSpillBytes)
-    {
-        this.maxSpillBytes = maxSpillBytes;
-        return this;
-    }
-
-    // TODO: All following configs are deprecated and will be removed after references from all
-    //  other systems are cleared.
-
-    @Config(REGISTER_TEST_FUNCTIONS)
-    public NativeExecutionSystemConfig setRegisterTestFunctions(boolean registerTestFunctions)
-    {
-        this.registerTestFunctions = registerTestFunctions;
-        return this;
-    }
-
-    public boolean isRegisterTestFunctions()
-    {
-        return registerTestFunctions;
-    }
-
-    @Config(MEMORY_POOL_INIT_CAPACITY)
-    public NativeExecutionSystemConfig setMemoryPoolInitCapacity(long memoryPoolInitCapacity)
-    {
-        this.memoryPoolInitCapacity = memoryPoolInitCapacity;
-        return this;
-    }
-
-    public long getMemoryPoolInitCapacity()
-    {
-        return memoryPoolInitCapacity;
-    }
-
-    @Config(MEMORY_POOL_RESERVED_CAPACITY)
-    public NativeExecutionSystemConfig setMemoryPoolReservedCapacity(long memoryPoolReservedCapacity)
-    {
-        this.memoryPoolReservedCapacity = memoryPoolReservedCapacity;
-        return this;
-    }
-
-    public long getMemoryPoolReservedCapacity()
+    public Map getAllProperties()
     {
-        return memoryPoolReservedCapacity;
+        ImmutableMap.Builder builder = ImmutableMap.builder();
+        builder.putAll(systemConfigs);
+        defaultSystemConfigs.entrySet().stream()
+                .filter(entry -> !systemConfigs.containsKey(entry.getKey()))
+                .forEach(entry -> builder.put(entry.getKey(), entry.getValue()));
+        return builder.build();
     }
 
-    @Config(MEMORY_POOL_TRANSFER_CAPACITY)
-    public NativeExecutionSystemConfig setMemoryPoolTransferCapacity(long memoryPoolTransferCapacity)
+    public NativeExecutionSystemConfig update(String key, String value)
     {
-        this.memoryPoolTransferCapacity = memoryPoolTransferCapacity;
+        systemConfigs.put(key, value);
         return this;
     }
 
-    public long getMemoryPoolTransferCapacity()
+    private String getAbsolutePath(String path)
     {
-        return memoryPoolTransferCapacity;
-    }
+        File absolutePath = new File(path);
+        if (!absolutePath.isAbsolute()) {
+            // In the case of SparkEnv is not initialed (e.g. unit test), we just use current location instead of calling SparkFiles.getRootDirectory() to avoid error.
+            String rootDirectory = SparkEnv$.MODULE$.get() != null ? SparkFiles.getRootDirectory() : ".";
+            absolutePath = new File(rootDirectory, path);
+        }
 
-    @Config(MEMORY_RECLAIM_WAIT_MS)
-    public NativeExecutionSystemConfig setMemoryReclaimWaitMs(long memoryReclaimWaitMs)
-    {
-        this.memoryReclaimWaitMs = memoryReclaimWaitMs;
-        return this;
-    }
+        if (!absolutePath.exists()) {
+            throw new IllegalArgumentException(format("File doesn't exist %s", absolutePath.getAbsolutePath()));
+        }
 
-    public long getMemoryReclaimWaitMs()
-    {
-        return memoryReclaimWaitMs;
+        return absolutePath.getAbsolutePath();
     }
 }
diff --git a/presto-spark-base/src/main/java/com/facebook/presto/spark/execution/property/PrestoSparkWorkerProperty.java b/presto-spark-base/src/main/java/com/facebook/presto/spark/execution/property/PrestoSparkWorkerProperty.java
index f83bcb1464e92..becc2c4b47d71 100644
--- a/presto-spark-base/src/main/java/com/facebook/presto/spark/execution/property/PrestoSparkWorkerProperty.java
+++ b/presto-spark-base/src/main/java/com/facebook/presto/spark/execution/property/PrestoSparkWorkerProperty.java
@@ -19,14 +19,14 @@
  * A utility class that helps with properties and its materialization.
  */
 public class PrestoSparkWorkerProperty
-        extends WorkerProperty
+        extends WorkerProperty
 {
     @Inject
     public PrestoSparkWorkerProperty(
-            NativeExecutionConnectorConfig connectorConfig,
+            NativeExecutionCatalogProperties catalogProperties,
             NativeExecutionNodeConfig nodeConfig,
             NativeExecutionSystemConfig systemConfig)
     {
-        super(connectorConfig, nodeConfig, systemConfig);
+        super(catalogProperties, nodeConfig, systemConfig);
     }
 }
diff --git a/presto-spark-base/src/main/java/com/facebook/presto/spark/execution/property/WorkerProperty.java b/presto-spark-base/src/main/java/com/facebook/presto/spark/execution/property/WorkerProperty.java
index 1a2faa3edf009..67b5aa22986fb 100644
--- a/presto-spark-base/src/main/java/com/facebook/presto/spark/execution/property/WorkerProperty.java
+++ b/presto-spark-base/src/main/java/com/facebook/presto/spark/execution/property/WorkerProperty.java
@@ -64,25 +64,25 @@
  * }
  * }
  */
-public class WorkerProperty
+public class WorkerProperty
 {
-    private final T1 connectorConfig;
+    private final T1 catalogProperties;
     private final T2 nodeConfig;
     private final T3 systemConfig;
 
     public WorkerProperty(
-            T1 connectorConfig,
+            T1 catalogProperties,
             T2 nodeConfig,
             T3 systemConfig)
     {
         this.systemConfig = requireNonNull(systemConfig, "systemConfig is null");
         this.nodeConfig = requireNonNull(nodeConfig, "nodeConfig is null");
-        this.connectorConfig = requireNonNull(connectorConfig, "connectorConfig is null");
+        this.catalogProperties = requireNonNull(catalogProperties, "catalogProperties is null");
     }
 
-    public T1 getConnectorConfig()
+    public T1 getCatalogProperties()
     {
-        return connectorConfig;
+        return catalogProperties;
     }
 
     public T2 getNodeConfig()
@@ -95,12 +95,12 @@ public T3 getSystemConfig()
         return systemConfig;
     }
 
-    public void populateAllProperties(Path systemConfigPath, Path nodeConfigPath, Path connectorConfigPath)
+    public void populateAllProperties(Path systemConfigPath, Path nodeConfigPath, Path catalogConfigsPath)
             throws IOException
     {
         populateProperty(systemConfig.getAllProperties(), systemConfigPath);
         populateProperty(nodeConfig.getAllProperties(), nodeConfigPath);
-        populateProperty(connectorConfig.getAllProperties(), connectorConfigPath);
+        populateCatalogProperties(catalogProperties.getAllCatalogProperties(), catalogConfigsPath);
     }
 
     private void populateProperty(Map properties, Path path)
@@ -122,4 +122,20 @@ private void populateProperty(Map properties, Path path)
             throw e;
         }
     }
+
+    private void populateCatalogProperties(Map> catalogProperties, Path path)
+            throws IOException
+    {
+        File catalogDir = path.toFile();
+        if (!catalogDir.exists()) {
+            catalogDir.mkdirs();
+        }
+
+        for (Map.Entry> catalogEntry : catalogProperties.entrySet()) {
+            String catalogName = catalogEntry.getKey();
+            Map properties = catalogEntry.getValue();
+            Path catalogConfigPath = path.resolve(catalogName + ".properties");
+            populateProperty(properties, catalogConfigPath);
+        }
+    }
 }
diff --git a/presto-spark-base/src/main/java/com/facebook/presto/spark/node/PrestoSparkInternalNodeManager.java b/presto-spark-base/src/main/java/com/facebook/presto/spark/node/PrestoSparkInternalNodeManager.java
index 1046a820ad367..7df3792f4b5f7 100644
--- a/presto-spark-base/src/main/java/com/facebook/presto/spark/node/PrestoSparkInternalNodeManager.java
+++ b/presto-spark-base/src/main/java/com/facebook/presto/spark/node/PrestoSparkInternalNodeManager.java
@@ -18,10 +18,12 @@
 import com.facebook.presto.metadata.InternalNode;
 import com.facebook.presto.metadata.InternalNodeManager;
 import com.facebook.presto.spi.ConnectorId;
+import com.facebook.presto.spi.NodeLoadMetrics;
 import com.facebook.presto.spi.NodeState;
 import com.google.common.collect.ImmutableSet;
 
 import java.net.URI;
+import java.util.Optional;
 import java.util.Set;
 import java.util.function.Consumer;
 
@@ -125,4 +127,10 @@ public void removeNodeChangeListener(Consumer listener)
     {
         throw new UnsupportedOperationException();
     }
+
+    @Override
+    public Optional getNodeLoadMetrics(String nodeIdentifier)
+    {
+        return Optional.empty();
+    }
 }
diff --git a/presto-spark-base/src/test/java/com/facebook/presto/spark/PrestoSparkQueryRunner.java b/presto-spark-base/src/test/java/com/facebook/presto/spark/PrestoSparkQueryRunner.java
index b31117343954b..0d5945fdb5a82 100644
--- a/presto-spark-base/src/test/java/com/facebook/presto/spark/PrestoSparkQueryRunner.java
+++ b/presto-spark-base/src/test/java/com/facebook/presto/spark/PrestoSparkQueryRunner.java
@@ -247,7 +247,8 @@ public static PrestoSparkQueryRunner createHivePrestoSparkQueryRunner(Iterable properties = process.getWorkerProperty().getSystemConfig().getAllProperties();
+        assertEquals(properties.get(NativeExecutionSystemConfig.SYSTEM_MEMORY_GB), "10");
+        assertEquals(properties.get(NativeExecutionSystemConfig.QUERY_MEMORY_GB), "8");
+        assertEquals(properties.get(NativeExecutionSystemConfig.QUERY_MAX_MEMORY_PER_NODE), "8GB");
+    }
+
+    @Test
+    public void testUpdateWorkerMemoryPropertiesWithOffHeapMemory()
+    {
+        // Test when spark.memory.offHeap.size is set to a value larger than current system memory
+        SparkConf sparkConf = new SparkConf();
+        sparkConf.set("spark.memory.offHeap.size", "20g"); // 20GB
+
+        TestingNativeExecutionProcess process = createTestingNativeExecutionProcess("10", "8", "8GB", sparkConf);
+        process.updateWorkerMemoryProperties();
+        // Verify the updated values
+        Map properties = process.getWorkerProperty().getSystemConfig().getAllProperties();
+
+        // Expected values:
+        // newSystemMemory = 20GB
+        // queryMemoryFraction = 8/10 = 0.8
+        // newQueryMemory = 20 * 0.8 = 16GB
+        assertEquals(properties.get(NativeExecutionSystemConfig.SYSTEM_MEMORY_GB), "20");
+        assertEquals(properties.get(NativeExecutionSystemConfig.QUERY_MEMORY_GB), "16");
+        assertEquals(properties.get(NativeExecutionSystemConfig.QUERY_MAX_MEMORY_PER_NODE), "16GB");
+    }
+
+    @Test
+    public void testUpdateWorkerMemoryPropertiesWithoutOffHeapSetting()
+    {
+        // Test when spark.memory.offHeap.size is not set
+        SparkConf sparkConf = new SparkConf();
+        // Don't set spark.memory.offHeap.size
+
+        TestingNativeExecutionProcess process = createTestingNativeExecutionProcess("10", "8", "8GB", sparkConf);
+        process.updateWorkerMemoryProperties();
+
+        // Verify that values remain unchanged
+        Map properties = process.getWorkerProperty().getSystemConfig().getAllProperties();
+        assertEquals(properties.get(NativeExecutionSystemConfig.SYSTEM_MEMORY_GB), "10");
+        assertEquals(properties.get(NativeExecutionSystemConfig.QUERY_MEMORY_GB), "8");
+        assertEquals(properties.get(NativeExecutionSystemConfig.QUERY_MAX_MEMORY_PER_NODE), "8GB");
+    }
+
+    @Test
+    public void testUpdateWorkerMemoryPropertiesWithZeroOffHeapMemory()
+    {
+        // Test when spark.memory.offHeap.size is set to 0
+        SparkConf sparkConf = new SparkConf();
+        sparkConf.set("spark.memory.offHeap.size", "0b");
+
+        TestingNativeExecutionProcess process = createTestingNativeExecutionProcess("10", "8", "8GB", sparkConf);
+        process.updateWorkerMemoryProperties();
+
+        // Verify that values remain unchanged when offHeapMemory is 0
+        Map properties = process.getWorkerProperty().getSystemConfig().getAllProperties();
+        assertEquals(properties.get(NativeExecutionSystemConfig.SYSTEM_MEMORY_GB), "10");
+        assertEquals(properties.get(NativeExecutionSystemConfig.QUERY_MEMORY_GB), "8");
+        assertEquals(properties.get(NativeExecutionSystemConfig.QUERY_MAX_MEMORY_PER_NODE), "8GB");
+    }
+
+    @Test
+    public void testUpdateWorkerMemoryPropertiesWithSmallerOffHeapMemory()
+    {
+        // Test when spark.memory.offHeap.size is smaller than current system memory
+        SparkConf sparkConf = new SparkConf();
+        sparkConf.set("spark.memory.offHeap.size", "5g"); // 5GB (smaller than current 10GB)
+
+        TestingNativeExecutionProcess process = createTestingNativeExecutionProcess("10", "8", "8GB", sparkConf);
+        process.updateWorkerMemoryProperties();
+
+        // Verify that values remain unchanged when offHeapMemory is smaller than current system memory
+        Map properties = process.getWorkerProperty().getSystemConfig().getAllProperties();
+        assertEquals(properties.get(NativeExecutionSystemConfig.SYSTEM_MEMORY_GB), "10");
+        assertEquals(properties.get(NativeExecutionSystemConfig.QUERY_MEMORY_GB), "8");
+        assertEquals(properties.get(NativeExecutionSystemConfig.QUERY_MAX_MEMORY_PER_NODE), "8GB");
+    }
+
+    @Test
+    public void testUpdateWorkerMemoryPropertiesWithDifferentQueryMemoryFraction()
+    {
+        // Test with different query memory to system memory ratio
+        SparkConf sparkConf = new SparkConf();
+        sparkConf.set("spark.memory.offHeap.size", "30g"); // 30GB
+
+        TestingNativeExecutionProcess process = createTestingNativeExecutionProcess("12", "6", "6GB", sparkConf);
+        process.updateWorkerMemoryProperties();
+
+        // Verify the updated values
+        Map properties = process.getWorkerProperty().getSystemConfig().getAllProperties();
+
+        // Expected values:
+        // newSystemMemory = 30GB
+        // queryMemoryFraction = 6/12 = 0.5
+        // newQueryMemory = 30 * 0.5 = 15GB
+        assertEquals(properties.get(NativeExecutionSystemConfig.SYSTEM_MEMORY_GB), "30");
+        assertEquals(properties.get(NativeExecutionSystemConfig.QUERY_MEMORY_GB), "15");
+        assertEquals(properties.get(NativeExecutionSystemConfig.QUERY_MAX_MEMORY_PER_NODE), "15GB");
+    }
+
+    @Test
+    public void testUpdateWorkerMemoryPropertiesWithZeroCurrentSystemMemory()
+    {
+        // Test edge case when current system memory is 0
+        SparkConf sparkConf = new SparkConf();
+        sparkConf.set("spark.memory.offHeap.size", "20g"); // 20GB
+
+        TestingNativeExecutionProcess process = createTestingNativeExecutionProcess("0", "8", "8GB", sparkConf);
+        process.updateWorkerMemoryProperties();
+
+        // Verify that values remain unchanged when current system memory is 0
+        Map properties = process.getWorkerProperty().getSystemConfig().getAllProperties();
+        assertEquals(properties.get(NativeExecutionSystemConfig.SYSTEM_MEMORY_GB), "0");
+        assertEquals(properties.get(NativeExecutionSystemConfig.QUERY_MEMORY_GB), "8");
+        assertEquals(properties.get(NativeExecutionSystemConfig.QUERY_MAX_MEMORY_PER_NODE), "8GB");
+    }
+
+    private TestingNativeExecutionProcess createTestingNativeExecutionProcess(
+            String systemMemoryGb,
+            String queryMemoryGb,
+            String queryMaxMemoryPerNode,
+            SparkConf sparkConf)
+    {
+        Map systemConfigs = new HashMap<>();
+        systemConfigs.put(NativeExecutionSystemConfig.SYSTEM_MEMORY_GB, systemMemoryGb);
+        systemConfigs.put(NativeExecutionSystemConfig.QUERY_MEMORY_GB, queryMemoryGb);
+        systemConfigs.put(NativeExecutionSystemConfig.QUERY_MAX_MEMORY_PER_NODE, queryMaxMemoryPerNode);
+
+        NativeExecutionSystemConfig systemConfig = new NativeExecutionSystemConfig(systemConfigs);
+        PrestoSparkWorkerProperty workerProperty = new PrestoSparkWorkerProperty(
+                new NativeExecutionCatalogProperties(ImmutableMap.of()),
+                new NativeExecutionNodeConfig(),
+                systemConfig);
+
+        Session session = testSessionBuilder().build();
+
+        try {
+            return new TestingNativeExecutionProcess(
+                    "/bin/echo",
+                    "",
+                    session,
+                    new TestPrestoSparkHttpClient.TestingOkHttpClient(newScheduledThreadPool(4),
+                            new TestPrestoSparkHttpClient.TestingResponseManager("test")),
+                    newSingleThreadExecutor(),
+                    newScheduledThreadPool(4),
+                    SERVER_INFO_JSON_CODEC,
+                    new Duration(10, TimeUnit.SECONDS),
+                    workerProperty,
+                    sparkConf);
+        }
+        catch (IOException e) {
+            throw new RuntimeException(e);
+        }
+    }
+
+    /**
+     * Testing subclass that allows injection of custom SparkConf for testing
+     */
+    private static class TestingNativeExecutionProcess
+            extends NativeExecutionProcess
+    {
+        private static final Logger log = Logger.get(TestingNativeExecutionProcess.class);
+        private final SparkConf testSparkConf;
+        public TestingNativeExecutionProcess(
+                String executablePath,
+                String programArguments,
+                Session session,
+                OkHttpClient httpClient,
+                java.util.concurrent.Executor executor,
+                ScheduledExecutorService scheduledExecutorService,
+                JsonCodec serverInfoCodec,
+                Duration maxErrorDuration,
+                PrestoSparkWorkerProperty workerProperty,
+                SparkConf testSparkConf)
+                throws IOException
+        {
+            super(executablePath, programArguments, session, httpClient, executor,
+                    scheduledExecutorService, serverInfoCodec, maxErrorDuration, workerProperty);
+            this.testSparkConf = testSparkConf;
+        }
+
+        @Override
+        protected SparkConf getSparkConf()
+        {
+            return testSparkConf;
+        }
+
+        public PrestoSparkWorkerProperty getWorkerProperty()
+        {
+            return super.getWorkerProperty();
+        }
+
+        @Override
+        protected void updateWorkerMemoryProperties()
+        {
+            super.updateWorkerMemoryProperties();
+        }
+    }
+
     private NativeExecutionProcessFactory createNativeExecutionProcessFactory()
     {
         TaskId taskId = new TaskId("testid", 0, 0, 0, 0);
         ScheduledExecutorService errorScheduler = newScheduledThreadPool(4);
         PrestoSparkWorkerProperty workerProperty = new PrestoSparkWorkerProperty(
-                new NativeExecutionConnectorConfig(),
+                new NativeExecutionCatalogProperties(ImmutableMap.of()),
                 new NativeExecutionNodeConfig(),
-                new NativeExecutionSystemConfig());
+                new NativeExecutionSystemConfig(ImmutableMap.of()));
         NativeExecutionProcessFactory factory = new NativeExecutionProcessFactory(
                 new TestPrestoSparkHttpClient.TestingOkHttpClient(
                         errorScheduler,
diff --git a/presto-spark-base/src/test/java/com/facebook/presto/spark/execution/http/TestPrestoSparkHttpClient.java b/presto-spark-base/src/test/java/com/facebook/presto/spark/execution/http/TestPrestoSparkHttpClient.java
index 2c0c3aa9ddb3e..700d55ade4e80 100644
--- a/presto-spark-base/src/test/java/com/facebook/presto/spark/execution/http/TestPrestoSparkHttpClient.java
+++ b/presto-spark-base/src/test/java/com/facebook/presto/spark/execution/http/TestPrestoSparkHttpClient.java
@@ -33,7 +33,7 @@
 import com.facebook.presto.spark.execution.nativeprocess.HttpNativeExecutionTaskResultFetcher;
 import com.facebook.presto.spark.execution.nativeprocess.NativeExecutionProcess;
 import com.facebook.presto.spark.execution.nativeprocess.NativeExecutionProcessFactory;
-import com.facebook.presto.spark.execution.property.NativeExecutionConnectorConfig;
+import com.facebook.presto.spark.execution.property.NativeExecutionCatalogProperties;
 import com.facebook.presto.spark.execution.property.NativeExecutionNodeConfig;
 import com.facebook.presto.spark.execution.property.NativeExecutionSystemConfig;
 import com.facebook.presto.spark.execution.property.PrestoSparkWorkerProperty;
@@ -897,9 +897,9 @@ private NativeExecutionProcess createNativeExecutionProcess(
             TestingResponseManager responseManager)
     {
         PrestoSparkWorkerProperty workerProperty = new PrestoSparkWorkerProperty(
-                new NativeExecutionConnectorConfig(),
+                new NativeExecutionCatalogProperties(ImmutableMap.of()),
                 new NativeExecutionNodeConfig(),
-                new NativeExecutionSystemConfig());
+                new NativeExecutionSystemConfig(ImmutableMap.of()));
         NativeExecutionProcessFactory factory = new NativeExecutionProcessFactory(
                 new TestingOkHttpClient(scheduledExecutorService, responseManager),
                 scheduledExecutorService,
diff --git a/presto-spark-base/src/test/java/com/facebook/presto/spark/execution/property/TestNativeExecutionSystemConfig.java b/presto-spark-base/src/test/java/com/facebook/presto/spark/execution/property/TestNativeExecutionSystemConfig.java
index 1298fcc48d831..0568193e7c815 100644
--- a/presto-spark-base/src/test/java/com/facebook/presto/spark/execution/property/TestNativeExecutionSystemConfig.java
+++ b/presto-spark-base/src/test/java/com/facebook/presto/spark/execution/property/TestNativeExecutionSystemConfig.java
@@ -14,9 +14,7 @@
 package com.facebook.presto.spark.execution.property;
 
 import com.facebook.airlift.configuration.testing.ConfigAssertions;
-import com.facebook.airlift.units.DataSize;
-import com.facebook.airlift.units.DataSize.Unit;
-import com.facebook.airlift.units.Duration;
+import com.google.common.collect.ImmutableMap;
 import org.testng.annotations.Test;
 
 import java.io.File;
@@ -27,7 +25,6 @@
 import java.nio.file.Paths;
 import java.util.Map;
 import java.util.Properties;
-import java.util.concurrent.TimeUnit;
 
 import static com.facebook.airlift.configuration.testing.ConfigAssertions.assertFullMapping;
 import static com.facebook.airlift.configuration.testing.ConfigAssertions.assertRecordedDefaults;
@@ -40,100 +37,145 @@ public class TestNativeExecutionSystemConfig
     public void testNativeExecutionSystemConfig()
     {
         // Test defaults
-        assertRecordedDefaults(ConfigAssertions.recordDefaults(NativeExecutionSystemConfig.class)
-                .setEnableSerializedPageChecksum(true)
-                .setEnableVeloxExpressionLogging(false)
-                .setEnableVeloxTaskLogging(true)
-                .setHttpServerReusePort(true)
-                .setHttpServerBindToNodeInternalAddressOnlyEnabled(true)
-                .setHttpServerPort(7777)
-                .setHttpServerNumIoThreadsHwMultiplier(1.0)
-                .setHttpsServerPort(7778)
-                .setEnableHttpsCommunication(false)
-                .setHttpsCiphers("AES128-SHA,AES128-SHA256,AES256-GCM-SHA384")
-                .setHttpsCertPath("")
-                .setHttpsKeyPath("")
-                .setExchangeHttpClientNumIoThreadsHwMultiplier(1.0)
-                .setAsyncDataCacheEnabled(false)
-                .setAsyncCacheSsdGb(0)
-                .setConnectorNumIoThreadsHwMultiplier(0.0)
-                .setShutdownOnsetSec(10)
-                .setSystemMemoryGb(10)
-                .setQueryMemoryGb(new DataSize(8, DataSize.Unit.GIGABYTE))
-                .setUseMmapAllocator(true)
-                .setMemoryArbitratorKind("SHARED")
-                .setMemoryArbitratorCapacityGb(8)
-                .setSharedArbitratorReservedCapacity(new DataSize(0, Unit.GIGABYTE))
-                .setSharedArbitratorMemoryPoolInitialCapacity(new DataSize(4, Unit.GIGABYTE))
-                .setSharedArbitratorMemoryPoolReservedCapacity(new DataSize(32, Unit.MEGABYTE))
-                .setSharedArbitratorMaxMemoryArbitrationTime(new Duration(5, TimeUnit.MINUTES))
-                .setSpillerSpillPath("")
-                .setConcurrentLifespansPerTask(5)
-                .setMaxDriversPerTask(15)
-                .setOldTaskCleanupMs(false)
-                .setPrestoVersion("dummy.presto.version")
-                .setShuffleName("local")
-                .setEnableHttpServerAccessLog(true)
-                .setCoreOnAllocationFailureEnabled(false)
-                .setSpillEnabled(true)
-                .setAggregationSpillEnabled(true)
-                .setJoinSpillEnabled(true)
-                .setOrderBySpillEnabled(true)
-                .setMaxSpillBytes(600L << 30)
-                .setRegisterTestFunctions(false)
-                .setMemoryPoolInitCapacity(0)
-                .setMemoryPoolReservedCapacity(0)
-                .setMemoryPoolTransferCapacity(0)
-                .setMemoryReclaimWaitMs(0));
+        NativeExecutionSystemConfig nativeExecutionSystemConfig = new NativeExecutionSystemConfig(
+                ImmutableMap.of());
+        ImmutableMap.Builder builder = ImmutableMap.builder();
+        ImmutableMap expectedConfigs = builder
+                .put("concurrent-lifespans-per-task", "5")
+                .put("enable-serialized-page-checksum", "true")
+                .put("enable_velox_expression_logging", "false")
+                .put("enable_velox_task_logging", "true")
+                .put("http-server.http.port", "7777")
+                .put("http-server.reuse-port", "true")
+                .put("http-server.bind-to-node-internal-address-only-enabled", "true")
+                .put("http-server.https.port", "7778")
+                .put("http-server.https.enabled", "false")
+                .put("https-supported-ciphers", "AES128-SHA,AES128-SHA256,AES256-GCM-SHA384")
+                .put("https-cert-path", "")
+                .put("https-key-path", "")
+                .put("http-server.num-io-threads-hw-multiplier", "1.0")
+                .put("exchange.http-client.num-io-threads-hw-multiplier", "1.0")
+                .put("async-data-cache-enabled", "false")
+                .put("async-cache-ssd-gb", "0")
+                .put("connector.num-io-threads-hw-multiplier", "0")
+                .put("presto.version", "dummy.presto.version")
+                .put("shutdown-onset-sec", "10")
+                .put("system-memory-gb", "10")
+                .put("query-memory-gb", "8")
+                .put("query.max-memory-per-node", "8GB")
+                .put("use-mmap-allocator", "true")
+                .put("memory-arbitrator-kind", "SHARED")
+                .put("shared-arbitrator.reserved-capacity", "0GB")
+                .put("shared-arbitrator.memory-pool-initial-capacity", "4GB")
+                .put("shared-arbitrator.max-memory-arbitration-time", "5m")
+                .put("experimental.spiller-spill-path", "")
+                .put("task.max-drivers-per-task", "15")
+                .put("enable-old-task-cleanup", "false")
+                .put("shuffle.name", "local")
+                .put("http-server.enable-access-log", "true")
+                .put("core-on-allocation-failure-enabled", "false")
+                .put("spill-enabled", "true")
+                .put("aggregation-spill-enabled", "true")
+                .put("join-spill-enabled", "true")
+                .put("order-by-spill-enabled", "true")
+                .put("max-spill-bytes", String.valueOf(600L << 30))
+                .build();
+        assertEquals(nativeExecutionSystemConfig.getAllProperties(), expectedConfigs);
 
         // Test explicit property mapping. Also makes sure properties returned by getAllProperties() covers full property list.
-        NativeExecutionSystemConfig expected = new NativeExecutionSystemConfig()
-                .setConcurrentLifespansPerTask(15)
-                .setEnableSerializedPageChecksum(false)
-                .setEnableVeloxExpressionLogging(true)
-                .setEnableVeloxTaskLogging(false)
-                .setHttpServerReusePort(false)
-                .setHttpServerBindToNodeInternalAddressOnlyEnabled(false)
-                .setHttpServerPort(8080)
-                .setHttpServerNumIoThreadsHwMultiplier(3.0)
-                .setHttpsServerPort(8081)
-                .setEnableHttpsCommunication(true)
-                .setHttpsCiphers("AES128-SHA")
-                .setHttpsCertPath("/tmp/non_existent.cert")
-                .setHttpsKeyPath("/tmp/non_existent.key")
-                .setExchangeHttpClientNumIoThreadsHwMultiplier(0.5)
-                .setAsyncDataCacheEnabled(true)
-                .setAsyncCacheSsdGb(1000)
-                .setConnectorNumIoThreadsHwMultiplier(1.0)
-                .setPrestoVersion("presto-version")
-                .setShutdownOnsetSec(30)
-                .setSystemMemoryGb(40)
-                .setQueryMemoryGb(new DataSize(20, Unit.GIGABYTE))
-                .setUseMmapAllocator(false)
-                .setMemoryArbitratorKind("")
-                .setMemoryArbitratorCapacityGb(10)
-                .setSharedArbitratorReservedCapacity(new DataSize(8, Unit.GIGABYTE))
-                .setSharedArbitratorMemoryPoolInitialCapacity(new DataSize(7, Unit.GIGABYTE))
-                .setSharedArbitratorMemoryPoolReservedCapacity(new DataSize(6, Unit.GIGABYTE))
-                .setSharedArbitratorMaxMemoryArbitrationTime(new Duration(123123123, TimeUnit.MILLISECONDS))
-                .setSpillerSpillPath("dummy.spill.path")
-                .setMaxDriversPerTask(30)
-                .setOldTaskCleanupMs(true)
-                .setShuffleName("custom")
-                .setEnableHttpServerAccessLog(false)
-                .setCoreOnAllocationFailureEnabled(true)
-                .setSpillEnabled(false)
-                .setAggregationSpillEnabled(false)
-                .setJoinSpillEnabled(false)
-                .setOrderBySpillEnabled(false)
-                .setMaxSpillBytes(1L)
-                .setRegisterTestFunctions(true)
-                .setMemoryPoolInitCapacity(10)
-                .setMemoryPoolReservedCapacity(10)
-                .setMemoryPoolTransferCapacity(10)
-                .setMemoryReclaimWaitMs(10);
-        Map properties = expected.getAllProperties();
-        assertFullMapping(properties, expected);
+        builder = ImmutableMap.builder();
+        ImmutableMap systemConfig = builder
+                .put("concurrent-lifespans-per-task", "15")
+                .put("enable-serialized-page-checksum", "false")
+                .put("enable_velox_expression_logging", "true")
+                .put("enable_velox_task_logging", "false")
+                .put("http-server.http.port", "8777")
+                .put("http-server.reuse-port", "false")
+                .put("http-server.bind-to-node-internal-address-only-enabled", "false")
+                .put("http-server.https.port", "8778")
+                .put("http-server.https.enabled", "true")
+                .put("https-supported-ciphers", "override-cipher")
+                .put("https-cert-path", "/override/path/cert")
+                .put("https-key-path", "/override/path/key")
+                .put("http-server.num-io-threads-hw-multiplier", "1.5")
+                .put("exchange.http-client.num-io-threads-hw-multiplier", "1.5")
+                .put("async-data-cache-enabled", "true")
+                .put("async-cache-ssd-gb", "1")
+                .put("connector.num-io-threads-hw-multiplier", "0.1")
+                .put("presto.version", "override.presto.version")
+                .put("shutdown-onset-sec", "15")
+                .put("system-memory-gb", "5")
+                .put("query-memory-gb", "4")
+                .put("query.max-memory-per-node", "4GB")
+                .put("use-mmap-allocator", "false")
+                .put("memory-arbitrator-kind", "NOOP")
+                .put("shared-arbitrator.reserved-capacity", "1GB")
+                .put("shared-arbitrator.memory-pool-initial-capacity", "1GB")
+                .put("shared-arbitrator.max-memory-arbitration-time", "1s")
+                .put("experimental.spiller-spill-path", "/abc")
+                .put("task.max-drivers-per-task", "25")
+                .put("enable-old-task-cleanup", "true")
+                .put("shuffle.name", "remote")
+                .put("http-server.enable-access-log", "false")
+                .put("core-on-allocation-failure-enabled", "true")
+                .put("spill-enabled", "false")
+                .put("aggregation-spill-enabled", "false")
+                .put("join-spill-enabled", "false")
+                .put("order-by-spill-enabled", "false")
+                .put("non-defined-property-key-0", "non-defined-property-value-0")
+                .put("non-defined-property-key-1", "non-defined-property-value-1")
+                .put("non-defined-property-key-2", "non-defined-property-value-2")
+                .put("non-defined-property-key-3", "non-defined-property-value-3")
+                .build();
+        nativeExecutionSystemConfig = new NativeExecutionSystemConfig(systemConfig);
+
+        builder = ImmutableMap.builder();
+        expectedConfigs = builder
+                .put("concurrent-lifespans-per-task", "15")
+                .put("enable-serialized-page-checksum", "false")
+                .put("enable_velox_expression_logging", "true")
+                .put("enable_velox_task_logging", "false")
+                .put("http-server.http.port", "8777")
+                .put("http-server.reuse-port", "false")
+                .put("http-server.bind-to-node-internal-address-only-enabled", "false")
+                .put("http-server.https.port", "8778")
+                .put("http-server.https.enabled", "true")
+                .put("https-supported-ciphers", "override-cipher")
+                .put("https-cert-path", "/override/path/cert")
+                .put("https-key-path", "/override/path/key")
+                .put("http-server.num-io-threads-hw-multiplier", "1.5")
+                .put("exchange.http-client.num-io-threads-hw-multiplier", "1.5")
+                .put("async-data-cache-enabled", "true")
+                .put("async-cache-ssd-gb", "1")
+                .put("connector.num-io-threads-hw-multiplier", "0.1")
+                .put("presto.version", "override.presto.version")
+                .put("shutdown-onset-sec", "15")
+                .put("system-memory-gb", "5")
+                .put("query-memory-gb", "4")
+                .put("query.max-memory-per-node", "4GB")
+                .put("use-mmap-allocator", "false")
+                .put("memory-arbitrator-kind", "NOOP")
+                .put("shared-arbitrator.reserved-capacity", "1GB")
+                .put("shared-arbitrator.memory-pool-initial-capacity", "1GB")
+                .put("shared-arbitrator.max-memory-arbitration-time", "1s")
+                .put("experimental.spiller-spill-path", "/abc")
+                .put("task.max-drivers-per-task", "25")
+                .put("enable-old-task-cleanup", "true")
+                .put("shuffle.name", "remote")
+                .put("http-server.enable-access-log", "false")
+                .put("core-on-allocation-failure-enabled", "true")
+                .put("spill-enabled", "false")
+                .put("aggregation-spill-enabled", "false")
+                .put("join-spill-enabled", "false")
+                .put("order-by-spill-enabled", "false")
+                .put("non-defined-property-key-0", "non-defined-property-value-0")
+                .put("non-defined-property-key-1", "non-defined-property-value-1")
+                .put("non-defined-property-key-2", "non-defined-property-value-2")
+                .put("non-defined-property-key-3", "non-defined-property-value-3")
+                .put("max-spill-bytes", String.valueOf(600L << 30)) // default spill bytes
+                .build();
+
+        assertEquals(nativeExecutionSystemConfig.getAllProperties(), expectedConfigs);
     }
 
     @Test
@@ -159,30 +201,26 @@ public void testNativeExecutionNodeConfig()
     }
 
     @Test
-    public void testNativeExecutionConnectorConfig()
+    public void testNativeExecutionCatalogProperties()
     {
-        // Test defaults
-        assertRecordedDefaults(ConfigAssertions.recordDefaults(NativeExecutionConnectorConfig.class)
-                .setCacheEnabled(false)
-                .setMaxCacheSize(new DataSize(0, DataSize.Unit.MEGABYTE))
-                .setConnectorName("hive"));
+        // Test default constructor
+        NativeExecutionCatalogProperties config = new NativeExecutionCatalogProperties(ImmutableMap.of());
+        assertEquals(config.getAllCatalogProperties(), ImmutableMap.of());
 
-        // Test explicit property mapping. Also makes sure properties returned by getAllProperties() covers full property list.
-        NativeExecutionConnectorConfig expected = new NativeExecutionConnectorConfig()
-                .setConnectorName("custom")
-                .setMaxCacheSize(new DataSize(32, DataSize.Unit.MEGABYTE))
-                .setCacheEnabled(true);
-        Map properties = new java.util.HashMap<>(expected.getAllProperties());
-        // Since the cache.max-cache-size requires to be size without the unit which to be compatible with the C++,
-        // here we convert the size from Long type (in string format) back to DataSize for comparison
-        properties.put("cache.max-cache-size", String.valueOf(new DataSize(Double.parseDouble(properties.get("cache.max-cache-size")), DataSize.Unit.MEGABYTE)));
-        assertFullMapping(properties, expected);
+        // Test constructor with catalog properties
+        Map> catalogProperties = ImmutableMap.of(
+                "hive", ImmutableMap.of("hive.metastore.uri", "thrift://localhost:9083"),
+                "tpch", ImmutableMap.of("tpch.splits-per-node", "4"));
+        NativeExecutionCatalogProperties configWithProps = new NativeExecutionCatalogProperties(catalogProperties);
+        assertEquals(configWithProps.getAllCatalogProperties(), catalogProperties);
     }
 
     @Test
     public void testFilePropertiesPopulator()
     {
-        PrestoSparkWorkerProperty workerProperty = new PrestoSparkWorkerProperty(new NativeExecutionConnectorConfig(), new NativeExecutionNodeConfig(), new NativeExecutionSystemConfig());
+        PrestoSparkWorkerProperty workerProperty = new PrestoSparkWorkerProperty(
+                new NativeExecutionCatalogProperties(ImmutableMap.of()), new NativeExecutionNodeConfig(),
+                new NativeExecutionSystemConfig(ImmutableMap.of()));
         testPropertiesPopulate(workerProperty);
     }
 
@@ -193,12 +231,17 @@ private void testPropertiesPopulate(PrestoSparkWorkerProperty workerProperty)
             directory = Files.createTempDirectory("presto");
             Path configPropertiesPath = Paths.get(directory.toString(), "config.properties");
             Path nodePropertiesPath = Paths.get(directory.toString(), "node.properties");
-            Path connectorPropertiesPath = Paths.get(directory.toString(), "catalog/hive.properties");
-            workerProperty.populateAllProperties(configPropertiesPath, nodePropertiesPath, connectorPropertiesPath);
+            Path catalogDirectory = Paths.get(directory.toString(), "catalog");  // Directory path for catalogs
+            workerProperty.populateAllProperties(configPropertiesPath, nodePropertiesPath, catalogDirectory);
 
             verifyProperties(workerProperty.getSystemConfig().getAllProperties(), readPropertiesFromDisk(configPropertiesPath));
             verifyProperties(workerProperty.getNodeConfig().getAllProperties(), readPropertiesFromDisk(nodePropertiesPath));
-            verifyProperties(workerProperty.getConnectorConfig().getAllProperties(), readPropertiesFromDisk(connectorPropertiesPath));
+            // Verify each catalog file was created properly
+            workerProperty.getCatalogProperties().getAllCatalogProperties().forEach(
+                    (catalogName, catalogProperties) -> {
+                        Path catalogFilePath = catalogDirectory.resolve(catalogName + ".properties");
+                        verifyProperties(catalogProperties, readPropertiesFromDisk(catalogFilePath));
+                    });
         }
         catch (Exception exception) {
             exception.printStackTrace();
diff --git a/presto-spark-classloader-interface/pom.xml b/presto-spark-classloader-interface/pom.xml
index 8f93be12818a4..a95e45827c5f1 100644
--- a/presto-spark-classloader-interface/pom.xml
+++ b/presto-spark-classloader-interface/pom.xml
@@ -13,7 +13,6 @@
     
         ${project.parent.basedir}
         true
-        2
     
 
     
@@ -23,11 +22,6 @@
             provided
         
 
-        
-            com.facebook.presto
-            presto-spark-classloader-spark${dep.pos.classloader.module-name.suffix}
-        
-
         
             com.google.guava
             guava
@@ -40,6 +34,25 @@
     
 
     
+        
+            spark2
+
+            
+                true
+                
+                    !spark-version
+                
+            
+
+            
+                
+                    com.facebook.presto
+                    presto-spark-classloader-spark2
+                    ${project.version}
+                
+            
+
+        
         
             spark3
 
@@ -50,22 +63,12 @@
                 
             
 
-            
-                3
-            
-
-            
-                
-                    
-                        com.facebook.presto.spark
-                        spark-core
-                        3.4.1-1
-                        compile
-                    
-                
-            
-
             
+                
+                    com.facebook.presto
+                    presto-spark-classloader-spark3
+                    ${project.version}
+                
                 
                     org.scala-lang
                     scala-library
diff --git a/presto-spark-package/pom.xml b/presto-spark-package/pom.xml
index 84ee9009f0bd7..44222dece8b37 100644
--- a/presto-spark-package/pom.xml
+++ b/presto-spark-package/pom.xml
@@ -152,7 +152,15 @@
 
         
             com.facebook.presto
-            presto-session-property-managers
+            presto-file-session-property-manager
+            ${project.version}
+            zip
+            provided
+        
+
+        
+            com.facebook.presto
+            presto-db-session-property-manager
             ${project.version}
             zip
             provided
diff --git a/presto-spark-package/src/main/assembly/presto.xml b/presto-spark-package/src/main/assembly/presto.xml
index 165bac55826e0..79446c79d59ba 100644
--- a/presto-spark-package/src/main/assembly/presto.xml
+++ b/presto-spark-package/src/main/assembly/presto.xml
@@ -85,8 +85,12 @@
             plugin/sqlserver
         
         
-            ${project.build.directory}/dependency/presto-session-property-managers-${project.version}
-            plugin/session-property-managers
+            ${project.build.directory}/dependency/presto-file-session-property-manager-${project.version}
+            plugin/file-session-property-manager
+        
+        
+            ${project.build.directory}/dependency/presto-db-session-property-manager-${project.version}
+            plugin/db-session-property-manager
         
     
 
diff --git a/presto-spark-testing/pom.xml b/presto-spark-testing/pom.xml
index 6ecb44edc0a25..14d2dcf21d548 100644
--- a/presto-spark-testing/pom.xml
+++ b/presto-spark-testing/pom.xml
@@ -88,7 +88,13 @@
 
         
             com.facebook.presto
-            presto-session-property-managers
+            presto-session-property-managers-common
+            test
+        
+
+        
+            com.facebook.presto
+            presto-file-session-property-manager
             test
         
 
diff --git a/presto-spark-testing/src/test/java/com/facebook/presto/spark/testing/TestPrestoSparkLauncherIntegrationSmokeTest.java b/presto-spark-testing/src/test/java/com/facebook/presto/spark/testing/TestPrestoSparkLauncherIntegrationSmokeTest.java
index 4d73d96734caf..e826d055e57e8 100644
--- a/presto-spark-testing/src/test/java/com/facebook/presto/spark/testing/TestPrestoSparkLauncherIntegrationSmokeTest.java
+++ b/presto-spark-testing/src/test/java/com/facebook/presto/spark/testing/TestPrestoSparkLauncherIntegrationSmokeTest.java
@@ -51,7 +51,7 @@
 import java.util.concurrent.TimeoutException;
 import java.util.regex.Pattern;
 
-import static com.facebook.presto.session.FileSessionPropertyManager.CODEC;
+import static com.facebook.presto.session.file.FileSessionPropertyManager.CODEC;
 import static com.facebook.presto.spark.testing.Processes.destroyProcess;
 import static com.facebook.presto.testing.TestingSession.testSessionBuilder;
 import static com.facebook.presto.tests.QueryAssertions.assertEqualsIgnoreOrder;
@@ -181,7 +181,8 @@ public void setUp()
                 Optional.of(Pattern.compile("global.*")),
                 Optional.empty(),
                 Optional.empty(),
-                properties);
+                properties,
+                ImmutableMap.of());
         sessionPropertyConfigJsonFile = new File(tempDir, "session-property-config.json");
         Files.write(sessionPropertyConfigJsonFile.toPath(), CODEC.toJsonBytes(Collections.singletonList(spec)));
         sessionPropertyConfig = new File(tempDir, "session-property-configuration.properties");
diff --git a/presto-spi/src/main/java/com/facebook/presto/spi/ConnectorTableLayout.java b/presto-spi/src/main/java/com/facebook/presto/spi/ConnectorTableLayout.java
index f355b52f7e993..d794088f14f79 100644
--- a/presto-spi/src/main/java/com/facebook/presto/spi/ConnectorTableLayout.java
+++ b/presto-spi/src/main/java/com/facebook/presto/spi/ConnectorTableLayout.java
@@ -34,6 +34,7 @@ public class ConnectorTableLayout
     private final Optional discretePredicates;
     private final List> localProperties;
     private final Optional remainingPredicate;
+    private final Optional uniqueColumn;
 
     public ConnectorTableLayout(ConnectorTableLayoutHandle handle)
     {
@@ -68,6 +69,20 @@ public ConnectorTableLayout(
             Optional discretePredicates,
             List> localProperties,
             Optional remainingPredicate)
+    {
+        this(handle, columns, predicate, tablePartitioning, streamPartitioningColumns, discretePredicates, localProperties, remainingPredicate, Optional.empty());
+    }
+
+    public ConnectorTableLayout(
+            ConnectorTableLayoutHandle handle,
+            Optional> columns,
+            TupleDomain predicate,
+            Optional tablePartitioning,
+            Optional> streamPartitioningColumns,
+            Optional discretePredicates,
+            List> localProperties,
+            Optional remainingPredicate,
+            Optional uniqueColumn)
     {
         requireNonNull(handle, "handle is null");
         requireNonNull(columns, "columns is null");
@@ -86,6 +101,7 @@ public ConnectorTableLayout(
         this.discretePredicates = discretePredicates;
         this.localProperties = localProperties;
         this.remainingPredicate = remainingPredicate;
+        this.uniqueColumn = uniqueColumn;
     }
 
     public ConnectorTableLayoutHandle getHandle()
@@ -123,6 +139,11 @@ public Optional getRemainingPredicate()
         return remainingPredicate;
     }
 
+    public Optional getUniqueColumn()
+    {
+        return uniqueColumn;
+    }
+
     /**
      * The partitioning of the table across the worker nodes.
      * 

diff --git a/presto-spi/src/main/java/com/facebook/presto/spi/NodeLoadMetrics.java b/presto-spi/src/main/java/com/facebook/presto/spi/NodeLoadMetrics.java
new file mode 100644
index 0000000000000..493018b344587
--- /dev/null
+++ b/presto-spi/src/main/java/com/facebook/presto/spi/NodeLoadMetrics.java
@@ -0,0 +1,123 @@
+/*
+ * Licensed 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 com.facebook.presto.spi;
+
+import com.facebook.drift.annotations.ThriftConstructor;
+import com.facebook.drift.annotations.ThriftField;
+import com.facebook.drift.annotations.ThriftStruct;
+import com.fasterxml.jackson.annotation.JsonCreator;
+import com.fasterxml.jackson.annotation.JsonProperty;
+
+import java.util.Objects;
+
+/**
+ * Load Statistics for a worker node in the Presto cluster along with the state.
+ */
+@ThriftStruct
+public class NodeLoadMetrics
+{
+    private final double cpuUsedPercent;
+    private final double memoryUsedInBytes;
+    private final int numQueuedDrivers;
+    private final boolean cpuOverload;
+    private final boolean memoryOverload;
+
+    @JsonCreator
+    @ThriftConstructor
+    public NodeLoadMetrics(
+            @JsonProperty("cpuUsedPercent") double cpuUsedPercent,
+            @JsonProperty("memoryUsedInBytes") double memoryUsedInBytes,
+            @JsonProperty("numQueuedDrivers") int numQueuedDrivers,
+            @JsonProperty("cpuOverload") boolean cpuOverload,
+            @JsonProperty("memoryOverload") boolean memoryOverload)
+    {
+        this.cpuUsedPercent = cpuUsedPercent;
+        this.memoryUsedInBytes = memoryUsedInBytes;
+        this.numQueuedDrivers = numQueuedDrivers;
+        this.cpuOverload = cpuOverload;
+        this.memoryOverload = memoryOverload;
+    }
+
+    @JsonProperty
+    @ThriftField(1)
+    public double getCpuUsedPercent()
+    {
+        return cpuUsedPercent;
+    }
+
+    @JsonProperty
+    @ThriftField(2)
+    public double getMemoryUsedInBytes()
+    {
+        return memoryUsedInBytes;
+    }
+
+    @JsonProperty
+    @ThriftField(3)
+    public int getNumQueuedDrivers()
+    {
+        return numQueuedDrivers;
+    }
+
+    @JsonProperty
+    @ThriftField(4)
+    public boolean getCpuOverload()
+    {
+        return cpuOverload;
+    }
+
+    @JsonProperty
+    @ThriftField(5)
+    public boolean getMemoryOverload()
+    {
+        return memoryOverload;
+    }
+
+    @Override
+    public boolean equals(Object o)
+    {
+        if (this == o) {
+            return true;
+        }
+        if (o == null || getClass() != o.getClass()) {
+            return false;
+        }
+        NodeLoadMetrics that = (NodeLoadMetrics) o;
+        return Double.compare(cpuUsedPercent, that.cpuUsedPercent) == 0 &&
+                Double.compare(memoryUsedInBytes, that.memoryUsedInBytes) == 0 &&
+                numQueuedDrivers == that.numQueuedDrivers &&
+                cpuOverload == that.cpuOverload &&
+                memoryOverload == that.memoryOverload;
+    }
+
+    @Override
+    public int hashCode()
+    {
+        return Objects.hash(cpuUsedPercent, memoryUsedInBytes, numQueuedDrivers, cpuOverload, memoryOverload);
+    }
+
+    @Override
+    public String toString()
+    {
+        StringBuilder sb = new StringBuilder();
+        sb.append("NodeLoadMetrics{");
+        sb.append("cpuUsedPercent=").append(cpuUsedPercent);
+        sb.append(", memoryUsedInBytes=").append(memoryUsedInBytes);
+        sb.append(", numQueuedDrivers=").append(numQueuedDrivers);
+        sb.append(", cpuOverload=").append(cpuOverload);
+        sb.append(", memoryOverload=").append(memoryOverload);
+        sb.append('}');
+        return sb.toString();
+    }
+}
diff --git a/presto-spi/src/main/java/com/facebook/presto/spi/NodeStats.java b/presto-spi/src/main/java/com/facebook/presto/spi/NodeStats.java
new file mode 100644
index 0000000000000..6eba1949d0199
--- /dev/null
+++ b/presto-spi/src/main/java/com/facebook/presto/spi/NodeStats.java
@@ -0,0 +1,90 @@
+/*
+ * Licensed 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 com.facebook.presto.spi;
+
+import com.facebook.drift.annotations.ThriftConstructor;
+import com.facebook.drift.annotations.ThriftField;
+import com.facebook.drift.annotations.ThriftStruct;
+import com.fasterxml.jackson.annotation.JsonCreator;
+import com.fasterxml.jackson.annotation.JsonProperty;
+
+import java.util.Objects;
+import java.util.Optional;
+
+import static java.util.Objects.requireNonNull;
+
+/**
+ * Load Statistics for a node in the Presto cluster along with the state.
+ */
+@ThriftStruct
+public class NodeStats
+{
+    private final NodeState nodeState;
+    private final Optional loadMetrics;
+
+    @JsonCreator
+    @ThriftConstructor
+    public NodeStats(
+            @JsonProperty("nodeState") NodeState nodeState,
+            @JsonProperty("loadMetrics") Optional loadMetrics)
+    {
+        this.nodeState = requireNonNull(nodeState, "nodeState is null");
+        this.loadMetrics = loadMetrics;
+    }
+
+    @JsonProperty
+    @ThriftField(1)
+    public NodeState getNodeState()
+    {
+        return nodeState;
+    }
+
+    @JsonProperty
+    @ThriftField(2)
+    public Optional getLoadMetrics()
+    {
+        return loadMetrics;
+    }
+
+    @Override
+    public boolean equals(Object o)
+    {
+        if (this == o) {
+            return true;
+        }
+        if (o == null || getClass() != o.getClass()) {
+            return false;
+        }
+        NodeStats nodeStats = (NodeStats) o;
+        return Objects.equals(nodeState, nodeStats.nodeState) &&
+                Objects.equals(loadMetrics, nodeStats.loadMetrics);
+    }
+
+    @Override
+    public int hashCode()
+    {
+        return Objects.hash(nodeState, loadMetrics);
+    }
+
+    @Override
+    public String toString()
+    {
+        StringBuilder sb = new StringBuilder();
+        sb.append("NodeStats{");
+        sb.append("nodeState=").append(nodeState);
+        sb.append(", loadMetrics=").append(loadMetrics);
+        sb.append('}');
+        return sb.toString();
+    }
+}
diff --git a/presto-spi/src/main/java/com/facebook/presto/spi/UniqueProperty.java b/presto-spi/src/main/java/com/facebook/presto/spi/UniqueProperty.java
new file mode 100644
index 0000000000000..946cdb6dc1dab
--- /dev/null
+++ b/presto-spi/src/main/java/com/facebook/presto/spi/UniqueProperty.java
@@ -0,0 +1,98 @@
+/*
+ * Licensed 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 com.facebook.presto.spi;
+
+import com.fasterxml.jackson.annotation.JsonCreator;
+import com.fasterxml.jackson.annotation.JsonProperty;
+
+import java.util.Collections;
+import java.util.Objects;
+import java.util.Optional;
+import java.util.Set;
+import java.util.function.Function;
+
+import static java.util.Objects.requireNonNull;
+
+public final class UniqueProperty
+        implements LocalProperty
+{
+    private final E column;
+
+    @JsonCreator
+    public UniqueProperty(@JsonProperty("column") E column)
+    {
+        this.column = requireNonNull(column, "column is null");
+    }
+
+    @Override
+    public boolean isOrderSensitive()
+    {
+        return false;
+    }
+
+    @JsonProperty
+    public E getColumn()
+    {
+        return column;
+    }
+
+    public Set getColumns()
+    {
+        return Collections.singleton(column);
+    }
+
+    @Override
+    public  Optional> translate(Function> translator)
+    {
+        return translator.apply(column)
+                .map(UniqueProperty::new);
+    }
+
+    @Override
+    public boolean isSimplifiedBy(LocalProperty known)
+    {
+        return known instanceof UniqueProperty && known.equals(this);
+    }
+
+    @Override
+    public Optional> withConstants(Set constants)
+    {
+        return Optional.of(this);
+    }
+
+    @Override
+    public String toString()
+    {
+        return "U(" + column + ")";
+    }
+
+    @Override
+    public boolean equals(Object o)
+    {
+        if (this == o) {
+            return true;
+        }
+        if (o == null || getClass() != o.getClass()) {
+            return false;
+        }
+        UniqueProperty that = (UniqueProperty) o;
+        return Objects.equals(column, that.column);
+    }
+
+    @Override
+    public int hashCode()
+    {
+        return Objects.hash(column);
+    }
+}
diff --git a/presto-spi/src/main/java/com/facebook/presto/spi/eventlistener/OutputColumnMetadata.java b/presto-spi/src/main/java/com/facebook/presto/spi/eventlistener/OutputColumnMetadata.java
new file mode 100644
index 0000000000000..d4016dd04f511
--- /dev/null
+++ b/presto-spi/src/main/java/com/facebook/presto/spi/eventlistener/OutputColumnMetadata.java
@@ -0,0 +1,77 @@
+/*
+ * Licensed 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 com.facebook.presto.spi.eventlistener;
+
+import com.facebook.presto.common.SourceColumn;
+import com.fasterxml.jackson.annotation.JsonCreator;
+import com.fasterxml.jackson.annotation.JsonProperty;
+
+import java.util.Objects;
+import java.util.Set;
+
+import static java.util.Objects.requireNonNull;
+
+public class OutputColumnMetadata
+{
+    private final String columnName;
+    private final String columnType;
+    private final Set sourceColumns;
+
+    @JsonCreator
+    public OutputColumnMetadata(String columnName, String columnType, Set sourceColumns)
+    {
+        this.columnName = requireNonNull(columnName, "columnName is null");
+        this.columnType = requireNonNull(columnType, "columnType is null");
+        this.sourceColumns = requireNonNull(sourceColumns, "sourceColumns is null");
+    }
+
+    @JsonProperty
+    public String getColumnName()
+    {
+        return columnName;
+    }
+
+    @JsonProperty
+    public String getColumnType()
+    {
+        return columnType;
+    }
+
+    @JsonProperty
+    public Set getSourceColumns()
+    {
+        return sourceColumns;
+    }
+
+    @Override
+    public int hashCode()
+    {
+        return Objects.hash(columnName, columnType, sourceColumns);
+    }
+
+    @Override
+    public boolean equals(Object obj)
+    {
+        if (this == obj) {
+            return true;
+        }
+        if ((obj == null) || (getClass() != obj.getClass())) {
+            return false;
+        }
+        OutputColumnMetadata other = (OutputColumnMetadata) obj;
+        return Objects.equals(columnName, other.columnName) &&
+                Objects.equals(columnType, other.columnType) &&
+                Objects.equals(sourceColumns, other.sourceColumns);
+    }
+}
diff --git a/presto-spi/src/main/java/com/facebook/presto/spi/eventlistener/QueryOutputMetadata.java b/presto-spi/src/main/java/com/facebook/presto/spi/eventlistener/QueryOutputMetadata.java
index 22b33bc695670..21acdbcf21407 100644
--- a/presto-spi/src/main/java/com/facebook/presto/spi/eventlistener/QueryOutputMetadata.java
+++ b/presto-spi/src/main/java/com/facebook/presto/spi/eventlistener/QueryOutputMetadata.java
@@ -30,7 +30,7 @@ public class QueryOutputMetadata
     private final Optional jsonLengthLimitExceeded;
 
     private final String serializedCommitOutput;
-    private final Optional> columns;
+    private final Optional> columns;
 
     public QueryOutputMetadata(
             String catalogName,
@@ -39,7 +39,7 @@ public QueryOutputMetadata(
             Optional connectorOutputMetadata,
             Optional jsonLengthLimitExceeded,
             String serializedCommitOutput,
-            Optional> columns)
+            Optional> columns)
     {
         this.catalogName = requireNonNull(catalogName, "catalogName is null");
         this.schema = requireNonNull(schema, "schema is null");
@@ -87,7 +87,7 @@ public String getSerializedCommitOutput()
     }
 
     @JsonProperty
-    public Optional> getColumns()
+    public Optional> getColumns()
     {
         return columns;
     }
diff --git a/presto-main-base/src/main/java/com/facebook/presto/type/LiteralParameter.java b/presto-spi/src/main/java/com/facebook/presto/spi/function/LiteralParameter.java
similarity index 95%
rename from presto-main-base/src/main/java/com/facebook/presto/type/LiteralParameter.java
rename to presto-spi/src/main/java/com/facebook/presto/spi/function/LiteralParameter.java
index 548419647d58b..f57ba796fdc53 100644
--- a/presto-main-base/src/main/java/com/facebook/presto/type/LiteralParameter.java
+++ b/presto-spi/src/main/java/com/facebook/presto/spi/function/LiteralParameter.java
@@ -11,7 +11,7 @@
  * See the License for the specific language governing permissions and
  * limitations under the License.
  */
-package com.facebook.presto.type;
+package com.facebook.presto.spi.function;
 
 import java.lang.annotation.ElementType;
 import java.lang.annotation.Retention;
diff --git a/presto-spi/src/main/java/com/facebook/presto/spi/function/table/ArgumentSpecification.java b/presto-spi/src/main/java/com/facebook/presto/spi/function/table/ArgumentSpecification.java
index 5443475c98ce0..73c822095863f 100644
--- a/presto-spi/src/main/java/com/facebook/presto/spi/function/table/ArgumentSpecification.java
+++ b/presto-spi/src/main/java/com/facebook/presto/spi/function/table/ArgumentSpecification.java
@@ -30,6 +30,7 @@
  */
 public abstract class ArgumentSpecification
 {
+    public static final String argumentType = "Abstract";
     private final String name;
     private final boolean required;
 
@@ -58,4 +59,9 @@ public Object getDefaultValue()
     {
         return defaultValue;
     }
+
+    public String getArgumentType()
+    {
+        return argumentType;
+    }
 }
diff --git a/presto-spi/src/main/java/com/facebook/presto/spi/function/table/DescriptorArgumentSpecification.java b/presto-spi/src/main/java/com/facebook/presto/spi/function/table/DescriptorArgumentSpecification.java
index 637d25fac73be..11ed93f02cd00 100644
--- a/presto-spi/src/main/java/com/facebook/presto/spi/function/table/DescriptorArgumentSpecification.java
+++ b/presto-spi/src/main/java/com/facebook/presto/spi/function/table/DescriptorArgumentSpecification.java
@@ -16,6 +16,7 @@
 public class DescriptorArgumentSpecification
         extends ArgumentSpecification
 {
+    public static final String argumentType = "DescriptorArgumentSpecification";
     private DescriptorArgumentSpecification(String name, boolean required, Descriptor defaultValue)
     {
         super(name, required, defaultValue);
@@ -52,4 +53,10 @@ public DescriptorArgumentSpecification build()
             return new DescriptorArgumentSpecification(name, required, defaultValue);
         }
     }
+
+    @Override
+    public String getArgumentType()
+    {
+        return argumentType;
+    }
 }
diff --git a/presto-spi/src/main/java/com/facebook/presto/spi/function/table/ReturnTypeSpecification.java b/presto-spi/src/main/java/com/facebook/presto/spi/function/table/ReturnTypeSpecification.java
index f93e9d2b601e3..2c77503cadada 100644
--- a/presto-spi/src/main/java/com/facebook/presto/spi/function/table/ReturnTypeSpecification.java
+++ b/presto-spi/src/main/java/com/facebook/presto/spi/function/table/ReturnTypeSpecification.java
@@ -27,6 +27,7 @@
  */
 public abstract class ReturnTypeSpecification
 {
+    public static final String returnType = "Abstract";
     /**
      * The proper columns of the table function are not known at function declaration time.
      * They must be determined at query analysis time based on the actual call arguments.
@@ -34,9 +35,16 @@ public abstract class ReturnTypeSpecification
     public static class GenericTable
             extends ReturnTypeSpecification
     {
+        public static final String returnType = "GenericTable";
         public static final GenericTable GENERIC_TABLE = new GenericTable();
 
         private GenericTable() {}
+
+        @Override
+        public String getReturnType()
+        {
+            return returnType;
+        }
     }
 
     /**
@@ -45,9 +53,16 @@ private GenericTable() {}
     public static class OnlyPassThrough
             extends ReturnTypeSpecification
     {
+        public static final String returnType = "OnlyPassThrough";
         public static final OnlyPassThrough ONLY_PASS_THROUGH = new OnlyPassThrough();
 
         private OnlyPassThrough() {}
+
+        @Override
+        public String getReturnType()
+        {
+            return returnType;
+        }
     }
 
     /**
@@ -57,6 +72,7 @@ private OnlyPassThrough() {}
     public static class DescribedTable
             extends ReturnTypeSpecification
     {
+        public static final String returnType = "DescribedTable";
         private final Descriptor descriptor;
 
         public DescribedTable(Descriptor descriptor)
@@ -70,5 +86,16 @@ public Descriptor getDescriptor()
         {
             return descriptor;
         }
+
+        @Override
+        public String getReturnType()
+        {
+            return returnType;
+        }
+    }
+
+    public String getReturnType()
+    {
+        return returnType;
     }
 }
diff --git a/presto-spi/src/main/java/com/facebook/presto/spi/function/table/ScalarArgumentSpecification.java b/presto-spi/src/main/java/com/facebook/presto/spi/function/table/ScalarArgumentSpecification.java
index 6016b4fc9b5ce..94f98bafe949a 100644
--- a/presto-spi/src/main/java/com/facebook/presto/spi/function/table/ScalarArgumentSpecification.java
+++ b/presto-spi/src/main/java/com/facebook/presto/spi/function/table/ScalarArgumentSpecification.java
@@ -23,6 +23,7 @@
 public class ScalarArgumentSpecification
         extends ArgumentSpecification
 {
+    public static final String argumentType = "ScalarArgumentSpecification";
     private final Type type;
 
     private ScalarArgumentSpecification(String name, Type type, boolean required, Object defaultValue)
@@ -77,4 +78,10 @@ public ScalarArgumentSpecification build()
             return new ScalarArgumentSpecification(name, type, required, defaultValue);
         }
     }
+
+    @Override
+    public String getArgumentType()
+    {
+        return argumentType;
+    }
 }
diff --git a/presto-spi/src/main/java/com/facebook/presto/spi/function/table/TableArgumentSpecification.java b/presto-spi/src/main/java/com/facebook/presto/spi/function/table/TableArgumentSpecification.java
index e1c1e638c6a9e..44c74258152dd 100644
--- a/presto-spi/src/main/java/com/facebook/presto/spi/function/table/TableArgumentSpecification.java
+++ b/presto-spi/src/main/java/com/facebook/presto/spi/function/table/TableArgumentSpecification.java
@@ -19,6 +19,7 @@
 public class TableArgumentSpecification
         extends ArgumentSpecification
 {
+    public static final String argumentType = "TableArgumentSpecification";
     private final boolean rowSemantics;
     private final boolean pruneWhenEmpty;
     private final boolean passThroughColumns;
@@ -107,4 +108,10 @@ public TableArgumentSpecification build()
             return new TableArgumentSpecification(name, rowSemantics, pruneWhenEmpty, passThroughColumns);
         }
     }
+
+    @Override
+    public String getArgumentType()
+    {
+        return argumentType;
+    }
 }
diff --git a/presto-spi/src/main/java/com/facebook/presto/spi/plan/TableWriterNode.java b/presto-spi/src/main/java/com/facebook/presto/spi/plan/TableWriterNode.java
index f5cd4aed80915..c13fdcf5fda4a 100644
--- a/presto-spi/src/main/java/com/facebook/presto/spi/plan/TableWriterNode.java
+++ b/presto-spi/src/main/java/com/facebook/presto/spi/plan/TableWriterNode.java
@@ -20,6 +20,7 @@
 import com.facebook.presto.spi.SchemaTableName;
 import com.facebook.presto.spi.SourceLocation;
 import com.facebook.presto.spi.TableHandle;
+import com.facebook.presto.spi.eventlistener.OutputColumnMetadata;
 import com.facebook.presto.spi.relation.VariableReferenceExpression;
 import com.fasterxml.jackson.annotation.JsonCreator;
 import com.fasterxml.jackson.annotation.JsonIgnore;
@@ -281,6 +282,8 @@ public abstract static class WriterTarget
 
         public abstract SchemaTableName getSchemaTableName();
 
+        public abstract Optional> getOutputColumns();
+
         @Override
         public abstract String toString();
     }
@@ -291,12 +294,14 @@ public static class CreateName
         private final ConnectorId connectorId;
         private final ConnectorTableMetadata tableMetadata;
         private final Optional layout;
+        private final Optional> columns;
 
-        public CreateName(ConnectorId connectorId, ConnectorTableMetadata tableMetadata, Optional layout)
+        public CreateName(ConnectorId connectorId, ConnectorTableMetadata tableMetadata, Optional layout, Optional> columns)
         {
             this.connectorId = requireNonNull(connectorId, "connectorId is null");
             this.tableMetadata = requireNonNull(tableMetadata, "tableMetadata is null");
             this.layout = requireNonNull(layout, "layout is null");
+            this.columns = requireNonNull(columns, "columns is null");
         }
 
         @Override
@@ -327,6 +332,12 @@ public String toString()
         {
             return connectorId + "." + tableMetadata.getTable();
         }
+
+        @Override
+        public Optional> getOutputColumns()
+        {
+            return columns;
+        }
     }
 
     public static class InsertReference
@@ -334,11 +345,13 @@ public static class InsertReference
     {
         private final TableHandle handle;
         private final SchemaTableName schemaTableName;
+        private final Optional> columns;
 
-        public InsertReference(TableHandle handle, SchemaTableName schemaTableName)
+        public InsertReference(TableHandle handle, SchemaTableName schemaTableName, Optional> columns)
         {
             this.handle = requireNonNull(handle, "handle is null");
             this.schemaTableName = requireNonNull(schemaTableName, "schemaTableName is null");
+            this.columns = requireNonNull(columns, "columns is null");
         }
 
         public TableHandle getHandle()
@@ -358,6 +371,12 @@ public SchemaTableName getSchemaTableName()
             return schemaTableName;
         }
 
+        @Override
+        public Optional> getOutputColumns()
+        {
+            return columns;
+        }
+
         @Override
         public String toString()
         {
@@ -396,6 +415,12 @@ public SchemaTableName getSchemaTableName()
             return schemaTableName;
         }
 
+        @Override
+        public Optional> getOutputColumns()
+        {
+            return Optional.empty();
+        }
+
         public String toString()
         {
             return handle.toString();
@@ -431,6 +456,12 @@ public SchemaTableName getSchemaTableName()
             return schemaTableName;
         }
 
+        @Override
+        public Optional> getOutputColumns()
+        {
+            return Optional.empty();
+        }
+
         @Override
         public String toString()
         {
@@ -474,6 +505,12 @@ public SchemaTableName getSchemaTableName()
             return schemaTableName;
         }
 
+        @Override
+        public Optional> getOutputColumns()
+        {
+            return Optional.empty();
+        }
+
         public List getUpdatedColumns()
         {
             return updatedColumns;
diff --git a/presto-main-base/src/main/java/com/facebook/presto/UnknownTypeException.java b/presto-spi/src/main/java/com/facebook/presto/spi/type/UnknownTypeException.java
similarity index 96%
rename from presto-main-base/src/main/java/com/facebook/presto/UnknownTypeException.java
rename to presto-spi/src/main/java/com/facebook/presto/spi/type/UnknownTypeException.java
index 984f953cbc0ee..9138efcf5d6ae 100644
--- a/presto-main-base/src/main/java/com/facebook/presto/UnknownTypeException.java
+++ b/presto-spi/src/main/java/com/facebook/presto/spi/type/UnknownTypeException.java
@@ -11,7 +11,7 @@
  * See the License for the specific language governing permissions and
  * limitations under the License.
  */
-package com.facebook.presto;
+package com.facebook.presto.spi.type;
 
 import com.facebook.presto.common.type.TypeSignature;
 import com.facebook.presto.spi.PrestoException;
diff --git a/presto-sql-invoked-functions-plugin/pom.xml b/presto-sql-invoked-functions-plugin/pom.xml
new file mode 100644
index 0000000000000..5ab338547b4ef
--- /dev/null
+++ b/presto-sql-invoked-functions-plugin/pom.xml
@@ -0,0 +1,29 @@
+
+    4.0.0
+    
+        com.facebook.presto
+        presto-root
+        0.295-SNAPSHOT
+    
+
+    presto-sql-invoked-functions-plugin
+    Presto - Sql invoked functions plugin
+    presto-plugin
+
+    
+        ${project.parent.basedir}
+    
+
+    
+        
+            com.facebook.presto
+            presto-spi
+            provided
+        
+        
+            com.google.guava
+            guava
+        
+    
+
diff --git a/presto-sql-invoked-functions-plugin/src/main/java/com/facebook/presto/scalar/sql/ArrayIntersectFunction.java b/presto-sql-invoked-functions-plugin/src/main/java/com/facebook/presto/scalar/sql/ArrayIntersectFunction.java
new file mode 100644
index 0000000000000..860faede7942d
--- /dev/null
+++ b/presto-sql-invoked-functions-plugin/src/main/java/com/facebook/presto/scalar/sql/ArrayIntersectFunction.java
@@ -0,0 +1,35 @@
+/*
+ * Licensed 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 com.facebook.presto.scalar.sql;
+
+import com.facebook.presto.spi.function.Description;
+import com.facebook.presto.spi.function.SqlInvokedScalarFunction;
+import com.facebook.presto.spi.function.SqlParameter;
+import com.facebook.presto.spi.function.SqlType;
+import com.facebook.presto.spi.function.TypeParameter;
+
+public class ArrayIntersectFunction
+{
+    private ArrayIntersectFunction() {}
+
+    @SqlInvokedScalarFunction(value = "array_intersect", deterministic = true, calledOnNullInput = false)
+    @Description("Intersects elements of all arrays in the given array")
+    @TypeParameter("T")
+    @SqlParameter(name = "input", type = "array>")
+    @SqlType("array")
+    public static String arrayIntersectArray()
+    {
+        return "RETURN reduce(input, IF((cardinality(input) = 0), ARRAY[], input[1]), (s, x) -> array_intersect(s, x), (s) -> s)";
+    }
+}
diff --git a/presto-main-base/src/main/java/com/facebook/presto/operator/scalar/sql/ArraySqlFunctions.java b/presto-sql-invoked-functions-plugin/src/main/java/com/facebook/presto/scalar/sql/ArraySqlFunctions.java
similarity index 99%
rename from presto-main-base/src/main/java/com/facebook/presto/operator/scalar/sql/ArraySqlFunctions.java
rename to presto-sql-invoked-functions-plugin/src/main/java/com/facebook/presto/scalar/sql/ArraySqlFunctions.java
index 2fc1218ea2408..b5ab56eb6f312 100644
--- a/presto-main-base/src/main/java/com/facebook/presto/operator/scalar/sql/ArraySqlFunctions.java
+++ b/presto-sql-invoked-functions-plugin/src/main/java/com/facebook/presto/scalar/sql/ArraySqlFunctions.java
@@ -11,7 +11,7 @@
  * See the License for the specific language governing permissions and
  * limitations under the License.
  */
-package com.facebook.presto.operator.scalar.sql;
+package com.facebook.presto.scalar.sql;
 
 import com.facebook.presto.spi.function.Description;
 import com.facebook.presto.spi.function.SqlInvokedScalarFunction;
diff --git a/presto-main-base/src/main/java/com/facebook/presto/operator/scalar/sql/MapNormalizeFunction.java b/presto-sql-invoked-functions-plugin/src/main/java/com/facebook/presto/scalar/sql/MapNormalizeFunction.java
similarity index 96%
rename from presto-main-base/src/main/java/com/facebook/presto/operator/scalar/sql/MapNormalizeFunction.java
rename to presto-sql-invoked-functions-plugin/src/main/java/com/facebook/presto/scalar/sql/MapNormalizeFunction.java
index 1618dac2c4356..99104f3027b8d 100644
--- a/presto-main-base/src/main/java/com/facebook/presto/operator/scalar/sql/MapNormalizeFunction.java
+++ b/presto-sql-invoked-functions-plugin/src/main/java/com/facebook/presto/scalar/sql/MapNormalizeFunction.java
@@ -11,7 +11,7 @@
  * See the License for the specific language governing permissions and
  * limitations under the License.
  */
-package com.facebook.presto.operator.scalar.sql;
+package com.facebook.presto.scalar.sql;
 
 import com.facebook.presto.spi.function.Description;
 import com.facebook.presto.spi.function.SqlInvokedScalarFunction;
diff --git a/presto-main-base/src/main/java/com/facebook/presto/operator/scalar/sql/MapSqlFunctions.java b/presto-sql-invoked-functions-plugin/src/main/java/com/facebook/presto/scalar/sql/MapSqlFunctions.java
similarity index 99%
rename from presto-main-base/src/main/java/com/facebook/presto/operator/scalar/sql/MapSqlFunctions.java
rename to presto-sql-invoked-functions-plugin/src/main/java/com/facebook/presto/scalar/sql/MapSqlFunctions.java
index 0cf9558a22989..498c068dedf39 100644
--- a/presto-main-base/src/main/java/com/facebook/presto/operator/scalar/sql/MapSqlFunctions.java
+++ b/presto-sql-invoked-functions-plugin/src/main/java/com/facebook/presto/scalar/sql/MapSqlFunctions.java
@@ -11,7 +11,7 @@
  * See the License for the specific language governing permissions and
  * limitations under the License.
  */
-package com.facebook.presto.operator.scalar.sql;
+package com.facebook.presto.scalar.sql;
 
 import com.facebook.presto.spi.function.Description;
 import com.facebook.presto.spi.function.SqlInvokedScalarFunction;
diff --git a/presto-main-base/src/main/java/com/facebook/presto/operator/scalar/sql/SimpleSamplingPercent.java b/presto-sql-invoked-functions-plugin/src/main/java/com/facebook/presto/scalar/sql/SimpleSamplingPercent.java
similarity index 96%
rename from presto-main-base/src/main/java/com/facebook/presto/operator/scalar/sql/SimpleSamplingPercent.java
rename to presto-sql-invoked-functions-plugin/src/main/java/com/facebook/presto/scalar/sql/SimpleSamplingPercent.java
index 473bbedafd637..5aa6b6b5d966e 100644
--- a/presto-main-base/src/main/java/com/facebook/presto/operator/scalar/sql/SimpleSamplingPercent.java
+++ b/presto-sql-invoked-functions-plugin/src/main/java/com/facebook/presto/scalar/sql/SimpleSamplingPercent.java
@@ -11,7 +11,7 @@
  * See the License for the specific language governing permissions and
  * limitations under the License.
  */
-package com.facebook.presto.operator.scalar.sql;
+package com.facebook.presto.scalar.sql;
 
 import com.facebook.presto.spi.function.Description;
 import com.facebook.presto.spi.function.SqlInvokedScalarFunction;
diff --git a/presto-sql-invoked-functions-plugin/src/main/java/com/facebook/presto/scalar/sql/SqlInvokedFunctionsPlugin.java b/presto-sql-invoked-functions-plugin/src/main/java/com/facebook/presto/scalar/sql/SqlInvokedFunctionsPlugin.java
new file mode 100644
index 0000000000000..9a1e3650f72fe
--- /dev/null
+++ b/presto-sql-invoked-functions-plugin/src/main/java/com/facebook/presto/scalar/sql/SqlInvokedFunctionsPlugin.java
@@ -0,0 +1,36 @@
+/*
+ * Licensed 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 com.facebook.presto.scalar.sql;
+
+import com.facebook.presto.spi.Plugin;
+import com.google.common.collect.ImmutableSet;
+
+import java.util.Set;
+
+public class SqlInvokedFunctionsPlugin
+        implements Plugin
+{
+    @Override
+    public Set> getSqlInvokedFunctions()
+    {
+        return ImmutableSet.>builder()
+                .add(ArraySqlFunctions.class)
+                .add(MapNormalizeFunction.class)
+                .add(MapSqlFunctions.class)
+                .add(SimpleSamplingPercent.class)
+                .add(StringSqlFunctions.class)
+                .add(ArrayIntersectFunction.class)
+                .build();
+    }
+}
diff --git a/presto-main-base/src/main/java/com/facebook/presto/operator/scalar/sql/StringSqlFunctions.java b/presto-sql-invoked-functions-plugin/src/main/java/com/facebook/presto/scalar/sql/StringSqlFunctions.java
similarity index 97%
rename from presto-main-base/src/main/java/com/facebook/presto/operator/scalar/sql/StringSqlFunctions.java
rename to presto-sql-invoked-functions-plugin/src/main/java/com/facebook/presto/scalar/sql/StringSqlFunctions.java
index 1b0e6e15661ed..9201e4bb1cf94 100644
--- a/presto-main-base/src/main/java/com/facebook/presto/operator/scalar/sql/StringSqlFunctions.java
+++ b/presto-sql-invoked-functions-plugin/src/main/java/com/facebook/presto/scalar/sql/StringSqlFunctions.java
@@ -12,7 +12,7 @@
  * limitations under the License.
  */
 
-package com.facebook.presto.operator.scalar.sql;
+package com.facebook.presto.scalar.sql;
 
 import com.facebook.presto.spi.function.Description;
 import com.facebook.presto.spi.function.SqlInvokedScalarFunction;
diff --git a/presto-teradata-functions/src/test/java/com/facebook/presto/teradata/functions/TestTeradataDateFunctions.java b/presto-teradata-functions/src/test/java/com/facebook/presto/teradata/functions/TestTeradataDateFunctions.java
index 9b4a72aacc792..16bd47af0d5c1 100644
--- a/presto-teradata-functions/src/test/java/com/facebook/presto/teradata/functions/TestTeradataDateFunctions.java
+++ b/presto-teradata-functions/src/test/java/com/facebook/presto/teradata/functions/TestTeradataDateFunctions.java
@@ -18,6 +18,8 @@
 import com.facebook.presto.common.type.SqlDate;
 import com.facebook.presto.common.type.TimestampType;
 import com.facebook.presto.operator.scalar.AbstractTestFunctions;
+import com.facebook.presto.sql.analyzer.FeaturesConfig;
+import com.facebook.presto.sql.analyzer.FunctionsConfig;
 import org.testng.annotations.BeforeClass;
 import org.testng.annotations.Test;
 
@@ -40,7 +42,7 @@ public class TestTeradataDateFunctions
 
     protected TestTeradataDateFunctions()
     {
-        super(SESSION);
+        super(SESSION, new FeaturesConfig(), new FunctionsConfig(), false);
     }
 
     @BeforeClass
diff --git a/presto-teradata-functions/src/test/java/com/facebook/presto/teradata/functions/TestTeradataFunctions.java b/presto-teradata-functions/src/test/java/com/facebook/presto/teradata/functions/TestTeradataFunctions.java
index be6efbffec1ef..e724ca94566e7 100644
--- a/presto-teradata-functions/src/test/java/com/facebook/presto/teradata/functions/TestTeradataFunctions.java
+++ b/presto-teradata-functions/src/test/java/com/facebook/presto/teradata/functions/TestTeradataFunctions.java
@@ -14,9 +14,12 @@
 package com.facebook.presto.teradata.functions;
 
 import com.facebook.presto.operator.scalar.AbstractTestFunctions;
+import com.facebook.presto.sql.analyzer.FeaturesConfig;
+import com.facebook.presto.sql.analyzer.FunctionsConfig;
 import org.testng.annotations.BeforeClass;
 import org.testng.annotations.Test;
 
+import static com.facebook.presto.SessionTestUtils.TEST_SESSION;
 import static com.facebook.presto.common.type.BigintType.BIGINT;
 import static com.facebook.presto.common.type.VarcharType.VARCHAR;
 import static com.facebook.presto.common.type.VarcharType.createVarcharType;
@@ -25,6 +28,11 @@
 public class TestTeradataFunctions
         extends AbstractTestFunctions
 {
+    public TestTeradataFunctions()
+    {
+        super(TEST_SESSION, new FeaturesConfig(), new FunctionsConfig(), false);
+    }
+
     @BeforeClass
     public void setUp()
     {
diff --git a/presto-test-coverage/pom.xml b/presto-test-coverage/pom.xml
index dc812c9675e32..849974018f171 100644
--- a/presto-test-coverage/pom.xml
+++ b/presto-test-coverage/pom.xml
@@ -480,7 +480,21 @@
 
         
             com.facebook.presto
-            presto-session-property-managers
+            presto-session-property-managers-common
+            ${project.version}
+            compile
+        
+
+        
+            com.facebook.presto
+            presto-db-session-property-manager
+            ${project.version}
+            compile
+        
+
+        
+            com.facebook.presto
+            presto-file-session-property-manager
             ${project.version}
             compile
         
diff --git a/presto-testing-server-launcher/pom.xml b/presto-testing-server-launcher/pom.xml
index 67acac907c506..dc51f0626e6a0 100644
--- a/presto-testing-server-launcher/pom.xml
+++ b/presto-testing-server-launcher/pom.xml
@@ -45,49 +45,58 @@
         
     
 
-    
-        
-            
-                org.apache.maven.plugins
-                maven-shade-plugin
-                
-                    
-                        package
-                        
-                            shade
-                        
+    
+        
+            executable-jar
+            
+                
+                    !skipExecutableJar
+                
+            
+            
+                
+                    
+                        org.apache.maven.plugins
+                        maven-shade-plugin
+                        
+                            
+                                package
+                                
+                                    shade
+                                
+                                
+                                    true
+                                    executable
+                                    
+                                        
+                                            
+                                                ${main-class}
+                                            
+                                        
+                                    
+                                
+                            
+                        
+                    
+
+                    
+                        org.skife.maven
+                        really-executable-jar-maven-plugin
                         
-                            true
-                            executable
-                            
-                                
-                                    
-                                        ${main-class}
-                                    
-                                
-                            
+                            -Xmx1G
+                            executable
                         
-                    
-                
-            
-
-            
-                org.skife.maven
-                really-executable-jar-maven-plugin
-                
-                    -Xmx1G
-                    executable
-                
-                
-                    
-                        package
-                        
-                            really-executable-jar
-                        
-                    
-                
-            
-        
-    
-
+                        
+                            
+                                package
+                                
+                                    really-executable-jar
+                                
+                            
+                        
+                    
+                
+            
+        
+    
 
diff --git a/presto-tests/pom.xml b/presto-tests/pom.xml
index 3623dd0991913..963a82a22adad 100644
--- a/presto-tests/pom.xml
+++ b/presto-tests/pom.xml
@@ -378,6 +378,13 @@
             ratis-metrics-default
             test
         
+
+        
+            com.facebook.presto
+            presto-sql-invoked-functions-plugin
+            ${project.version}
+            test
+        
     
 
     
diff --git a/presto-tests/src/main/java/com/facebook/presto/tests/AbstractTestDistributedQueries.java b/presto-tests/src/main/java/com/facebook/presto/tests/AbstractTestDistributedQueries.java
index aaeebab0b6bf1..48ba3ecfc11fa 100644
--- a/presto-tests/src/main/java/com/facebook/presto/tests/AbstractTestDistributedQueries.java
+++ b/presto-tests/src/main/java/com/facebook/presto/tests/AbstractTestDistributedQueries.java
@@ -77,7 +77,6 @@
 import static com.facebook.presto.testing.TestingSession.TESTING_CATALOG;
 import static com.facebook.presto.testing.assertions.Assert.assertEquals;
 import static com.facebook.presto.tests.QueryAssertions.assertContains;
-import static com.facebook.presto.transaction.TransactionBuilder.transaction;
 import static com.google.common.collect.ImmutableList.toImmutableList;
 import static com.google.common.collect.Iterables.getOnlyElement;
 import static com.google.common.util.concurrent.Uninterruptibles.sleepUninterruptibly;
@@ -249,38 +248,72 @@ public void testCreateTable()
         assertFalse(getQueryRunner().tableExists(getSession(), "test_create_like"));
     }
 
+    @Test
+    public void testNonAutoCommitTransactionWithRollback()
+    {
+        assertUpdate("create table multi_statements_transaction_rollback(a int, b varchar)");
+        assertUpdate("insert into multi_statements_transaction_rollback values(1, '1001'), (2, '1002')", 2);
+
+        Session session = assertStartTransaction(getSession(), "start transaction");
+
+        assertQuery(session, "select * from multi_statements_transaction_rollback", "values(1, '1001'), (2, '1002')");
+        assertUpdate(session, "insert into multi_statements_transaction_rollback values(3, '1003'), (4, '1004')", 2);
+
+        // `Rollback` executes successfully, and the client gets a flag `clearTransactionId = true`
+        session = assertEndTransaction(session, "rollback");
+
+        assertQuery(session, "select * from multi_statements_transaction_rollback", "values(1, '1001'), (2, '1002')");
+        assertUpdate(session, "drop table if exists multi_statements_transaction_rollback");
+    }
+
+    @Test
+    public void testNonAutoCommitTransactionWithCommit()
+    {
+        assertUpdate("create table multi_statements_transaction_commit(a int, b varchar)");
+        assertUpdate("insert into multi_statements_transaction_commit values(1, '1001'), (2, '1002')", 2);
+        Session session = assertStartTransaction(getSession(), "start transaction");
+
+        assertQuery(session, "select * from multi_statements_transaction_commit", "values(1, '1001'), (2, '1002')");
+        assertUpdate(session, "insert into multi_statements_transaction_commit values(3, '1003'), (4, '1004')", 2);
+
+        // `Commit` executes successfully, and the client gets a flag `clearTransactionId = true`
+        session = assertEndTransaction(session, "commit");
+
+        assertQuery(session, "select * from multi_statements_transaction_commit",
+                "values(1, '1001'), (2, '1002'), (3, '1003'), (4, '1004')");
+        assertUpdate(session, "drop table if exists multi_statements_transaction_commit");
+    }
+
     @Test
     public void testNonAutoCommitTransactionWithFailAndRollback()
     {
         assertUpdate("create table test_non_autocommit_table(a int, b varchar)");
-        Session session = getQueryRunner().getDefaultSession();
-        String defaultCatalog = session.getCatalog().get();
-        transaction(getQueryRunner().getTransactionManager(), getQueryRunner().getAccessControl())
-                .execute(Session.builder(session)
-                                .setIdentity(new Identity("admin",
-                                        Optional.empty(),
-                                        ImmutableMap.of(defaultCatalog, new SelectedRole(ROLE, Optional.of("admin"))),
-                                        ImmutableMap.of(),
-                                        ImmutableMap.of(),
-                                        Optional.empty(),
-                                        Optional.empty()))
-                                .build(),
-                        txnSession -> {
-                            // simulate failure of SQL statement execution
-                            assertQueryFails(txnSession, "SELECT fail('forced failure')", "forced failure");
-
-                            // cannot execute any SQLs except `rollback` in current session
-                            assertQueryFails(txnSession, "select count(*) from test_non_autocommit_table", "Current transaction is aborted, commands ignored until end of transaction block");
-                            assertQueryFails(txnSession, "show tables", "Current transaction is aborted, commands ignored until end of transaction block");
-                            assertQueryFails(txnSession, "insert into test_non_autocommit_table values(1, '1001')", "Current transaction is aborted, commands ignored until end of transaction block");
-                            assertQueryFails(txnSession, "create table test_table(a int, b varchar)", "Current transaction is aborted, commands ignored until end of transaction block");
-
-                            // execute `rollback` successfully
-                            assertUpdate(txnSession, "rollback");
-                        });
-
-        assertQuery("select count(*) from test_non_autocommit_table", "values(0)");
-        assertUpdate("drop table if exists test_non_autocommit_table");
+        Session session = Session.builder(getSession())
+                .setIdentity(new Identity("admin",
+                        Optional.empty(),
+                        ImmutableMap.of(getSession().getCatalog().get(), new SelectedRole(ROLE, Optional.of("admin"))),
+                        ImmutableMap.of(),
+                        ImmutableMap.of(),
+                        Optional.empty(),
+                        Optional.empty()))
+                .build();
+
+        session = assertStartTransaction(session, "start transaction");
+
+        // simulate failure of SQL statement execution
+        assertQueryFails(session, "SELECT fail('forced failure')", "forced failure");
+
+        // cannot execute any SQLs except `rollback` in current session
+        assertQueryFails(session, "select count(*) from test_non_autocommit_table", "Current transaction is aborted, commands ignored until end of transaction block");
+        assertQueryFails(session, "show tables", "Current transaction is aborted, commands ignored until end of transaction block");
+        assertQueryFails(session, "insert into test_non_autocommit_table values(1, '1001')", "Current transaction is aborted, commands ignored until end of transaction block");
+        assertQueryFails(session, "create table test_table(a int, b varchar)", "Current transaction is aborted, commands ignored until end of transaction block");
+
+        // execute `rollback` successfully
+        session = assertEndTransaction(session, "rollback");
+
+        assertQuery(session, "select count(*) from test_non_autocommit_table", "values(0)");
+        assertUpdate(session, "drop table if exists test_non_autocommit_table");
     }
 
     @Test
diff --git a/presto-tests/src/main/java/com/facebook/presto/tests/AbstractTestQueries.java b/presto-tests/src/main/java/com/facebook/presto/tests/AbstractTestQueries.java
index b35829e2760a8..52774715e6bdd 100644
--- a/presto-tests/src/main/java/com/facebook/presto/tests/AbstractTestQueries.java
+++ b/presto-tests/src/main/java/com/facebook/presto/tests/AbstractTestQueries.java
@@ -3218,7 +3218,8 @@ public void testSetSessionNativeWorkerSessionProperty()
                 "MaterializedResult{rows=[[true]], " +
                         "types=[boolean], " +
                         "setSessionProperties={native_expression_max_array_size_in_reduce=50000}, " +
-                        "resetSessionProperties=[], updateType=SET SESSION}");
+                        "resetSessionProperties=[], updateType=SET SESSION, " +
+                        "clearTransactionId=false}");
     }
 
     @Test
diff --git a/presto-tests/src/main/java/com/facebook/presto/tests/AbstractTestQueryFramework.java b/presto-tests/src/main/java/com/facebook/presto/tests/AbstractTestQueryFramework.java
index 72999fbe97c16..9115059ca9e6d 100644
--- a/presto-tests/src/main/java/com/facebook/presto/tests/AbstractTestQueryFramework.java
+++ b/presto-tests/src/main/java/com/facebook/presto/tests/AbstractTestQueryFramework.java
@@ -268,6 +268,16 @@ protected void assertUpdate(@Language("SQL") String sql)
         assertUpdate(getSession(), sql);
     }
 
+    protected Session assertStartTransaction(Session session, @Language("SQL") String sql)
+    {
+        return QueryAssertions.assertStartTransaction(queryRunner, session, sql);
+    }
+
+    protected Session assertEndTransaction(Session session, @Language("SQL") String sql)
+    {
+        return QueryAssertions.assertEndTransaction(queryRunner, session, sql);
+    }
+
     protected void assertUpdate(Session session, @Language("SQL") String sql)
     {
         QueryAssertions.assertUpdate(queryRunner, session, sql, OptionalLong.empty(), Optional.empty());
@@ -565,7 +575,7 @@ protected SubPlan subplan(String sql, Session session)
         }
     }
 
-    private QueryExplainer getQueryExplainer()
+    protected QueryExplainer getQueryExplainer()
     {
         Metadata metadata = queryRunner.getMetadata();
         FeaturesConfig featuresConfig = createFeaturesConfig();
@@ -630,6 +640,12 @@ protected ExpectedQueryRunner getExpectedQueryRunner()
         return expectedQueryRunner;
     }
 
+    protected SqlParser getSqlParser()
+    {
+        checkState(sqlParser != null, "sqlParser not set");
+        return sqlParser;
+    }
+
     public interface QueryRunnerSupplier
     {
         QueryRunner get()
diff --git a/presto-tests/src/main/java/com/facebook/presto/tests/AbstractTestingPrestoClient.java b/presto-tests/src/main/java/com/facebook/presto/tests/AbstractTestingPrestoClient.java
index 9eb4d3eda142b..96d812c802e7f 100644
--- a/presto-tests/src/main/java/com/facebook/presto/tests/AbstractTestingPrestoClient.java
+++ b/presto-tests/src/main/java/com/facebook/presto/tests/AbstractTestingPrestoClient.java
@@ -117,7 +117,7 @@ public ResultWithQueryId execute(Session session, @Language("SQL") String sql
 
                 resultsSession.setWarnings(results.getWarnings());
 
-                T result = resultsSession.build(client.getSetSessionProperties(), client.getResetSessionProperties());
+                T result = resultsSession.build(client.getSetSessionProperties(), client.getResetSessionProperties(), client.getStartedTransactionId(), client.isClearTransactionId());
                 return new ResultWithQueryId<>(new QueryId(results.getId()), result);
             }
 
diff --git a/presto-tests/src/main/java/com/facebook/presto/tests/DistributedQueryRunner.java b/presto-tests/src/main/java/com/facebook/presto/tests/DistributedQueryRunner.java
index d9ce3439011ab..6949fc4e9c8b3 100644
--- a/presto-tests/src/main/java/com/facebook/presto/tests/DistributedQueryRunner.java
+++ b/presto-tests/src/main/java/com/facebook/presto/tests/DistributedQueryRunner.java
@@ -1040,6 +1040,21 @@ public void loadPlanCheckerProviderManager(String planCheckerProviderName, Map> planAssertion)
     {
         long start = System.nanoTime();
diff --git a/presto-tests/src/main/java/com/facebook/presto/tests/ResultsSession.java b/presto-tests/src/main/java/com/facebook/presto/tests/ResultsSession.java
index e37f28ba2dc09..2be0ce2629c0e 100644
--- a/presto-tests/src/main/java/com/facebook/presto/tests/ResultsSession.java
+++ b/presto-tests/src/main/java/com/facebook/presto/tests/ResultsSession.java
@@ -16,6 +16,7 @@
 import com.facebook.presto.client.QueryData;
 import com.facebook.presto.client.QueryStatusInfo;
 import com.facebook.presto.spi.PrestoWarning;
+import jakarta.annotation.Nullable;
 
 import java.util.List;
 import java.util.Map;
@@ -40,5 +41,5 @@ default void setWarnings(List warnings)
 
     void addResults(QueryStatusInfo statusInfo, QueryData data);
 
-    T build(Map setSessionProperties, Set resetSessionProperties);
+    T build(Map setSessionProperties, Set resetSessionProperties, @Nullable String startTransactionId, boolean clearTransactionId);
 }
diff --git a/presto-tests/src/main/java/com/facebook/presto/tests/TestTableFunctionInvocation.java b/presto-tests/src/main/java/com/facebook/presto/tests/TestTableFunctionInvocation.java
new file mode 100644
index 0000000000000..067c750f09497
--- /dev/null
+++ b/presto-tests/src/main/java/com/facebook/presto/tests/TestTableFunctionInvocation.java
@@ -0,0 +1,94 @@
+/*
+ * Licensed 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 com.facebook.presto.tests;
+
+import com.facebook.presto.connector.tvf.TestTVFConnectorColumnHandle;
+import com.facebook.presto.connector.tvf.TestTVFConnectorFactory;
+import com.facebook.presto.connector.tvf.TestTVFConnectorPlugin;
+import com.facebook.presto.connector.tvf.TestingTableFunctions.SimpleTableFunction;
+import com.facebook.presto.connector.tvf.TestingTableFunctions.SimpleTableFunction.SimpleTableFunctionHandle;
+import com.facebook.presto.spi.ConnectorSession;
+import com.facebook.presto.spi.ConnectorTableHandle;
+import com.facebook.presto.spi.connector.TableFunctionApplicationResult;
+import com.facebook.presto.testing.QueryRunner;
+import com.google.common.collect.ImmutableSet;
+import org.testng.annotations.BeforeClass;
+import org.testng.annotations.Test;
+
+import java.util.Map;
+import java.util.Optional;
+import java.util.function.BiFunction;
+import java.util.stream.IntStream;
+
+import static com.facebook.presto.common.type.VarcharType.createUnboundedVarcharType;
+import static com.facebook.presto.testing.TestingSession.testSessionBuilder;
+import static com.google.common.collect.ImmutableMap.toImmutableMap;
+
+public class TestTableFunctionInvocation
+        extends AbstractTestQueryFramework
+{
+    private static final String TESTING_CATALOG = "testing_catalog1";
+    private static final String TABLE_FUNCTION_SCHEMA = "table_function_schema";
+
+    @Override
+    protected QueryRunner createQueryRunner()
+            throws Exception
+    {
+        return DistributedQueryRunner.builder(testSessionBuilder()
+                        .setCatalog(TESTING_CATALOG)
+                        .setSchema(TABLE_FUNCTION_SCHEMA)
+                        .build())
+                .build();
+    }
+
+    @BeforeClass
+    public void setUp()
+    {
+        DistributedQueryRunner queryRunner = getDistributedQueryRunner();
+
+        BiFunction> getColumnHandles = (session, tableHandle) -> IntStream.range(0, 100)
+                .boxed()
+                .map(i -> "column_" + i)
+                .collect(toImmutableMap(column -> column, column -> new TestTVFConnectorColumnHandle(column, createUnboundedVarcharType()) {}));
+
+        queryRunner.installPlugin(new TestTVFConnectorPlugin(TestTVFConnectorFactory.builder()
+                .withTableFunctions(ImmutableSet.of(new SimpleTableFunction()))
+                .withApplyTableFunction((session, handle) -> {
+                    if (handle instanceof SimpleTableFunctionHandle) {
+                        SimpleTableFunctionHandle functionHandle = (SimpleTableFunctionHandle) handle;
+                        return Optional.of(new TableFunctionApplicationResult<>(functionHandle.getTableHandle(), functionHandle.getTableHandle().getColumns().orElseThrow(() -> new IllegalStateException("Columns are missing"))));
+                    }
+                    throw new IllegalStateException("Unsupported table function handle: " + handle.getClass().getSimpleName());
+                }).withGetColumnHandles(getColumnHandles)
+                .build()));
+        queryRunner.createCatalog(TESTING_CATALOG, "testTVF");
+    }
+
+    @Test
+    public void testPrimitiveDefaultArgument()
+    {
+        assertQuery("SELECT boolean_column FROM TABLE(system.simple_table_function(column => 'boolean_column', ignored => 1))", "SELECT true WHERE false");
+
+        // skip the `ignored` argument.
+        assertQuery("SELECT boolean_column FROM TABLE(system.simple_table_function(column => 'boolean_column'))",
+                "SELECT true WHERE false");
+    }
+
+    @Test
+    public void testNoArgumentsPassed()
+    {
+        assertQuery("SELECT col FROM TABLE(system.simple_table_function())",
+                  "SELECT true WHERE false");
+    }
+}
diff --git a/presto-tests/src/main/java/com/facebook/presto/tests/TestingPrestoClient.java b/presto-tests/src/main/java/com/facebook/presto/tests/TestingPrestoClient.java
index 475143e1e543b..26ddac7e426e5 100644
--- a/presto-tests/src/main/java/com/facebook/presto/tests/TestingPrestoClient.java
+++ b/presto-tests/src/main/java/com/facebook/presto/tests/TestingPrestoClient.java
@@ -18,6 +18,7 @@
 import com.facebook.presto.client.IntervalYearMonth;
 import com.facebook.presto.client.QueryData;
 import com.facebook.presto.client.QueryStatusInfo;
+import com.facebook.presto.common.transaction.TransactionId;
 import com.facebook.presto.common.type.ArrayType;
 import com.facebook.presto.common.type.BigintEnumType;
 import com.facebook.presto.common.type.DecimalType;
@@ -41,6 +42,7 @@
 import com.facebook.presto.type.SqlIntervalYearMonth;
 import com.google.common.base.Function;
 import com.google.common.collect.ImmutableList;
+import jakarta.annotation.Nullable;
 
 import java.math.BigDecimal;
 import java.time.LocalDate;
@@ -147,7 +149,7 @@ public void addResults(QueryStatusInfo statusInfo, QueryData data)
         }
 
         @Override
-        public MaterializedResult build(Map setSessionProperties, Set resetSessionProperties)
+        public MaterializedResult build(Map setSessionProperties, Set resetSessionProperties, @Nullable String startTransactionId, boolean clearTransactionId)
         {
             checkState(types.get() != null, "never received types for the query");
             return new MaterializedResult(
@@ -157,6 +159,8 @@ public MaterializedResult build(Map setSessionProperties, Set> getSqlInvokedFunctions()
                     .add(TestDuplicateSqlInvokedFunctions.class)
                     .build();
         }
+
+        @Override
+        public Set> getFunctions()
+        {
+            return ImmutableSet.>builder()
+                    .add(TestFunctions.class)
+                    // Adding a SQL Invoked function in the built-in functions to mimic a conflict.
+                    .add(ArrayIntersectFunction.class)
+                    .build();
+        }
     }
 
+    // As soon as we trigger the conflict check with the built-in functions, an error will be thrown if duplicate signatures are found.
+    // In `PrestoServer.java` this conflict check is triggered implicitly.
     @Test
     public void testDuplicateFunctionsLoaded()
     {
-        assertInvalidFunction(JAVA_BUILTIN_NAMESPACE + ".modulo(10,3)",
-                Pattern.quote(format("java.lang.IllegalArgumentException: Function already registered: %s.array_intersect(array(array(T))):array(T)", JAVA_BUILTIN_NAMESPACE)));
+        assertThrows(IllegalArgumentException.class, () -> server.triggerConflictCheckWithBuiltInFunctions());
     }
 }
diff --git a/presto-tests/src/test/java/com/facebook/presto/server/TestServerInfoResource.java b/presto-tests/src/test/java/com/facebook/presto/server/TestServerInfoResource.java
index 75f743d5e710c..f4130226f75f7 100644
--- a/presto-tests/src/test/java/com/facebook/presto/server/TestServerInfoResource.java
+++ b/presto-tests/src/test/java/com/facebook/presto/server/TestServerInfoResource.java
@@ -22,10 +22,10 @@
 import com.facebook.drift.transport.netty.codec.Protocol;
 import com.facebook.presto.server.testing.TestingPrestoServer;
 import com.facebook.presto.spi.NodeState;
+import com.facebook.presto.spi.NodeStats;
 import com.facebook.presto.tests.DistributedQueryRunner;
 import com.google.common.collect.ImmutableMap;
 import org.testng.annotations.AfterClass;
-import org.testng.annotations.AfterGroups;
 import org.testng.annotations.BeforeClass;
 import org.testng.annotations.BeforeGroups;
 import org.testng.annotations.DataProvider;
@@ -44,12 +44,14 @@
 import static com.facebook.presto.tests.tpch.TpchQueryRunner.createQueryRunnerWithNoClusterReadyCheck;
 import static org.testng.Assert.assertEquals;
 
-@Test(singleThreaded = true)
+@Test(singleThreaded = true, groups = "TestServerInfoResource")
 public class TestServerInfoResource
 {
     private HttpClient client;
-    private DistributedQueryRunner queryRunner;
-    private DistributedQueryRunner queryRunnerWithNoClusterReadyCheck;
+    private DistributedQueryRunner queryRunnerActive;
+    private DistributedQueryRunner queryRunnerInactiveResourceManagers;
+    private DistributedQueryRunner queryRunnerInactiveCoordinators;
+    private DistributedQueryRunner queryRunnerInactiveResourceGroup;
     private ThriftCodecManager thriftCodeManager;
 
     @BeforeClass
@@ -73,11 +75,28 @@ public void teardown()
         this.client = null;
     }
 
-    @AfterGroups(groups = {"createQueryRunner", "getServerStateWithoutRequiredResourceManagers", "getServerStateWithoutRequiredCoordinators", "getServerStateWithoutRequiredCoordinators", "createQueryRunnerWithNoClusterReadyCheckSkipLoadingResourceGroupConfigurationManager"})
+    @AfterClass(alwaysRun = true)
     public void serverTearDown()
     {
-        for (TestingPrestoServer server : queryRunner.getServers()) {
-            closeQuietly(server);
+        if (queryRunnerActive != null) {
+            for (TestingPrestoServer server : queryRunnerActive.getServers()) {
+                closeQuietly(server);
+            }
+        }
+        if (queryRunnerInactiveResourceManagers != null) {
+            for (TestingPrestoServer server : queryRunnerInactiveResourceManagers.getServers()) {
+                closeQuietly(server);
+            }
+        }
+        if (queryRunnerInactiveCoordinators != null) {
+            for (TestingPrestoServer server : queryRunnerInactiveCoordinators.getServers()) {
+                closeQuietly(server);
+            }
+        }
+        if (queryRunnerInactiveResourceGroup != null) {
+            for (TestingPrestoServer server : queryRunnerInactiveResourceGroup.getServers()) {
+                closeQuietly(server);
+            }
         }
     }
 
@@ -85,7 +104,7 @@ public void serverTearDown()
     public void createQueryRunnerSetup()
             throws Exception
     {
-        queryRunner = createQueryRunner(
+        queryRunnerActive = createQueryRunner(
                 ImmutableMap.of(),
                 ImmutableMap.of(),
                 ImmutableMap.of(),
@@ -96,17 +115,27 @@ public void createQueryRunnerSetup()
     @Test(timeOut = 30_000, groups = {"createQueryRunner"}, dataProvider = "thriftEncodingToggle")
     public void testGetServerStateWithRequiredResourceManagerCoordinators(boolean useThriftEncoding, Protocol thriftProtocol)
     {
-        TestingPrestoServer server = queryRunner.getCoordinator(0);
+        TestingPrestoServer server = queryRunnerActive.getCoordinator(0);
         URI uri = uriBuilderFrom(server.getBaseUrl().resolve("/v1/info/state")).build();
         NodeState state = getNodeState(uri, useThriftEncoding, thriftProtocol);
         assertEquals(state, NodeState.ACTIVE);
     }
 
+    @Test(timeOut = 30_000, groups = {"createQueryRunner"}, dataProvider = "thriftEncodingToggle", dependsOnMethods = "testGetServerStateWithRequiredResourceManagerCoordinators")
+    public void testGetServerStatsWithRequiredResourceManagerCoordinators(boolean useThriftEncoding, Protocol thriftProtocol)
+    {
+        TestingPrestoServer server = queryRunnerActive.getCoordinator(0);
+        URI uri = uriBuilderFrom(server.getBaseUrl().resolve("/v1/info/stats")).build();
+        NodeStats stats = getNodeStats(uri, useThriftEncoding, thriftProtocol);
+        assertEquals(stats.getNodeState(), NodeState.ACTIVE);
+        assertEquals(stats.getLoadMetrics().isPresent(), false);
+    }
+
     @BeforeGroups("getServerStateWithoutRequiredResourceManagers")
     public void createQueryRunnerWithNoClusterReadyCheckSetup()
             throws Exception
     {
-        queryRunner = createQueryRunnerWithNoClusterReadyCheck(
+        queryRunnerInactiveResourceManagers = createQueryRunnerWithNoClusterReadyCheck(
                         ImmutableMap.of(),
                         ImmutableMap.of(),
                         ImmutableMap.of(),
@@ -117,17 +146,27 @@ public void createQueryRunnerWithNoClusterReadyCheckSetup()
     @Test(timeOut = 30_000, groups = {"getServerStateWithoutRequiredResourceManagers"}, dataProvider = "thriftEncodingToggle")
     public void testGetServerStateWithoutRequiredResourceManagers(boolean useThriftEncoding, Protocol thriftProtocol)
     {
-        TestingPrestoServer server = queryRunner.getCoordinator(0);
+        TestingPrestoServer server = queryRunnerInactiveResourceManagers.getCoordinator(0);
         URI uri = uriBuilderFrom(server.getBaseUrl().resolve("/v1/info/state")).build();
         NodeState state = getNodeState(uri, useThriftEncoding, thriftProtocol);
         assertEquals(state, NodeState.INACTIVE);
     }
 
+    @Test(timeOut = 30_000, groups = {"getServerStateWithoutRequiredResourceManagers"}, dataProvider = "thriftEncodingToggle", dependsOnMethods = "testGetServerStateWithoutRequiredResourceManagers")
+    public void testGetServerStatsWithoutRequiredResourceManagers(boolean useThriftEncoding, Protocol thriftProtocol)
+    {
+        TestingPrestoServer server = queryRunnerInactiveResourceManagers.getCoordinator(0);
+        URI uri = uriBuilderFrom(server.getBaseUrl().resolve("/v1/info/stats")).build();
+        NodeStats stats = getNodeStats(uri, useThriftEncoding, thriftProtocol);
+        assertEquals(stats.getNodeState(), NodeState.INACTIVE);
+        assertEquals(stats.getLoadMetrics().isPresent(), false);
+    }
+
     @BeforeGroups("getServerStateWithoutRequiredCoordinators")
     public void getServerStateWithoutRequiredCoordinatorsSetup()
             throws Exception
     {
-        queryRunner = createQueryRunnerWithNoClusterReadyCheck(
+        queryRunnerInactiveCoordinators = createQueryRunnerWithNoClusterReadyCheck(
                 ImmutableMap.of(),
                 ImmutableMap.of(),
                 ImmutableMap.of(),
@@ -138,18 +177,28 @@ public void getServerStateWithoutRequiredCoordinatorsSetup()
     @Test(timeOut = 30_000, groups = {"getServerStateWithoutRequiredCoordinators"}, dataProvider = "thriftEncodingToggle")
     public void testGetServerStateWithoutRequiredCoordinators(boolean useThriftEncoding, Protocol thriftProtocol)
     {
-        TestingPrestoServer server = queryRunner.getCoordinator(0);
+        TestingPrestoServer server = queryRunnerInactiveCoordinators.getCoordinator(0);
         URI uri = uriBuilderFrom(server.getBaseUrl().resolve("/v1/info/state")).build();
         NodeState state = getNodeState(uri, useThriftEncoding, thriftProtocol);
 
         assertEquals(state, NodeState.INACTIVE);
     }
 
+    @Test(timeOut = 30_000, groups = {"getServerStateWithoutRequiredCoordinators"}, dataProvider = "thriftEncodingToggle", dependsOnMethods = "testGetServerStateWithoutRequiredCoordinators")
+    public void testGetServerStatsWithoutRequiredCoordinators(boolean useThriftEncoding, Protocol thriftProtocol)
+    {
+        TestingPrestoServer server = queryRunnerInactiveCoordinators.getCoordinator(0);
+        URI uri = uriBuilderFrom(server.getBaseUrl().resolve("/v1/info/stats")).build();
+        NodeStats stats = getNodeStats(uri, useThriftEncoding, thriftProtocol);
+        assertEquals(stats.getNodeState(), NodeState.INACTIVE);
+        assertEquals(stats.getLoadMetrics().isPresent(), false);
+    }
+
     @BeforeGroups("createQueryRunnerWithNoClusterReadyCheckSkipLoadingResourceGroupConfigurationManager")
     public void createQueryRunnerWithNoClusterReadyCheckSkipLoadingResourceGroupConfigurationManager()
             throws Exception
     {
-        queryRunner = createQueryRunnerWithNoClusterReadyCheck(
+        queryRunnerInactiveResourceGroup = createQueryRunnerWithNoClusterReadyCheck(
                 ImmutableMap.of(),
                 ImmutableMap.of(),
                 ImmutableMap.of(),
@@ -161,12 +210,22 @@ public void createQueryRunnerWithNoClusterReadyCheckSkipLoadingResourceGroupConf
     public void testGetServerStateWhenResourceGroupConfigurationManagerNotLoaded()
             throws Exception
     {
-        TestingPrestoServer server = queryRunner.getCoordinator(0);
+        TestingPrestoServer server = queryRunnerInactiveResourceGroup.getCoordinator(0);
         URI uri = uriBuilderFrom(server.getBaseUrl().resolve("/v1/info/state")).build();
         NodeState state = getNodeState(uri, false, null);
         assertEquals(state, NodeState.INACTIVE);
     }
 
+    @Test(groups = {"createQueryRunnerWithNoClusterReadyCheckSkipLoadingResourceGroupConfigurationManager"}, dependsOnMethods = "testGetServerStateWhenResourceGroupConfigurationManagerNotLoaded")
+    public void testGetServerStatsWhenResourceGroupConfigurationManagerNotLoaded()
+    {
+        TestingPrestoServer server = queryRunnerInactiveResourceGroup.getCoordinator(0);
+        URI uri = uriBuilderFrom(server.getBaseUrl().resolve("/v1/info/stats")).build();
+        NodeStats stats = getNodeStats(uri, false, null);
+        assertEquals(stats.getNodeState(), NodeState.INACTIVE);
+        assertEquals(stats.getLoadMetrics().isPresent(), false);
+    }
+
     private NodeState getNodeState(URI uri, boolean useThriftEncoding, Protocol thriftProtocol)
     {
         Request.Builder requestBuilder = useThriftEncoding ? ThriftRequestUtils.prepareThriftGet(thriftProtocol) : getJsonTransportBuilder(prepareGet());
@@ -181,4 +240,19 @@ private NodeState getNodeState(URI uri, boolean useThriftEncoding, Protocol thri
             return client.execute(request, createJsonResponseHandler(jsonCodec(NodeState.class)));
         }
     }
+
+    private NodeStats getNodeStats(URI uri, boolean useThriftEncoding, Protocol thriftProtocol)
+    {
+        Request.Builder requestBuilder = useThriftEncoding ? ThriftRequestUtils.prepareThriftGet(thriftProtocol) : getJsonTransportBuilder(prepareGet());
+        Request request = requestBuilder
+                .setHeader(PRESTO_USER, "user")
+                .setUri(uri)
+                .build();
+        if (useThriftEncoding) {
+            return client.execute(request, new ThriftResponseHandler<>(thriftCodeManager.getCodec(NodeStats.class))).getValue();
+        }
+        else {
+            return client.execute(request, createJsonResponseHandler(jsonCodec(NodeStats.class)));
+        }
+    }
 }
diff --git a/presto-tests/src/test/java/com/facebook/presto/tests/TestDistributedEngineOnlyQueries.java b/presto-tests/src/test/java/com/facebook/presto/tests/TestDistributedEngineOnlyQueries.java
index b18173ff0bf20..8d6910ccaa59d 100644
--- a/presto-tests/src/test/java/com/facebook/presto/tests/TestDistributedEngineOnlyQueries.java
+++ b/presto-tests/src/test/java/com/facebook/presto/tests/TestDistributedEngineOnlyQueries.java
@@ -13,6 +13,7 @@
  */
 package com.facebook.presto.tests;
 
+import com.facebook.presto.scalar.sql.SqlInvokedFunctionsPlugin;
 import com.facebook.presto.testing.QueryRunner;
 import com.facebook.presto.tests.tpch.TpchQueryRunnerBuilder;
 
@@ -23,6 +24,8 @@ public class TestDistributedEngineOnlyQueries
     protected QueryRunner createQueryRunner()
             throws Exception
     {
-        return TpchQueryRunnerBuilder.builder().build();
+        QueryRunner queryRunner = TpchQueryRunnerBuilder.builder().build();
+        queryRunner.installPlugin(new SqlInvokedFunctionsPlugin());
+        return queryRunner;
     }
 }
diff --git a/presto-tests/src/test/java/com/facebook/presto/tests/TestDistributedSpilledQueries.java b/presto-tests/src/test/java/com/facebook/presto/tests/TestDistributedSpilledQueries.java
index e945493c2969e..1ab8522eb50a3 100644
--- a/presto-tests/src/test/java/com/facebook/presto/tests/TestDistributedSpilledQueries.java
+++ b/presto-tests/src/test/java/com/facebook/presto/tests/TestDistributedSpilledQueries.java
@@ -15,6 +15,7 @@
 
 import com.facebook.presto.Session;
 import com.facebook.presto.SystemSessionProperties;
+import com.facebook.presto.scalar.sql.SqlInvokedFunctionsPlugin;
 import com.facebook.presto.testing.QueryRunner;
 import com.facebook.presto.tpch.TpchPlugin;
 import com.google.common.collect.ImmutableMap;
@@ -63,6 +64,7 @@ public static QueryRunner localCreateQueryRunner()
         try {
             queryRunner.installPlugin(new TpchPlugin());
             queryRunner.createCatalog("tpch", "tpch");
+            queryRunner.installPlugin(new SqlInvokedFunctionsPlugin());
             return queryRunner;
         }
         catch (Exception e) {
diff --git a/presto-tests/src/test/java/com/facebook/presto/tests/TestDistributedSpilledQueriesWithTempStorage.java b/presto-tests/src/test/java/com/facebook/presto/tests/TestDistributedSpilledQueriesWithTempStorage.java
index c4ca56a3f78c0..739c5dc826631 100644
--- a/presto-tests/src/test/java/com/facebook/presto/tests/TestDistributedSpilledQueriesWithTempStorage.java
+++ b/presto-tests/src/test/java/com/facebook/presto/tests/TestDistributedSpilledQueriesWithTempStorage.java
@@ -15,6 +15,7 @@
 
 import com.facebook.presto.Session;
 import com.facebook.presto.SystemSessionProperties;
+import com.facebook.presto.scalar.sql.SqlInvokedFunctionsPlugin;
 import com.facebook.presto.testing.QueryRunner;
 import com.facebook.presto.tpch.TpchPlugin;
 import com.google.common.collect.ImmutableMap;
@@ -57,6 +58,8 @@ public static DistributedQueryRunner localCreateQueryRunner()
         try {
             queryRunner.installPlugin(new TpchPlugin());
             queryRunner.createCatalog("tpch", "tpch");
+            queryRunner.installPlugin(new SqlInvokedFunctionsPlugin());
+
             return queryRunner;
         }
         catch (Exception e) {
diff --git a/presto-tests/src/test/java/com/facebook/presto/tests/TestLocalQueries.java b/presto-tests/src/test/java/com/facebook/presto/tests/TestLocalQueries.java
index b381b01f5ada3..e7f2cf5a9a781 100644
--- a/presto-tests/src/test/java/com/facebook/presto/tests/TestLocalQueries.java
+++ b/presto-tests/src/test/java/com/facebook/presto/tests/TestLocalQueries.java
@@ -15,6 +15,7 @@
 
 import com.facebook.presto.Session;
 import com.facebook.presto.metadata.SessionPropertyManager;
+import com.facebook.presto.scalar.sql.SqlInvokedFunctionsPlugin;
 import com.facebook.presto.spi.CatalogSchemaTableName;
 import com.facebook.presto.spi.ConnectorId;
 import com.facebook.presto.sql.planner.planPrinter.IOPlanPrinter.ColumnConstraint;
@@ -78,6 +79,8 @@ public static LocalQueryRunner createLocalQueryRunner()
         sessionPropertyManager.addSystemSessionProperties(TEST_SYSTEM_PROPERTIES);
         sessionPropertyManager.addConnectorSessionProperties(new ConnectorId(TESTING_CATALOG), TEST_CATALOG_PROPERTIES);
 
+        localQueryRunner.installPlugin(new SqlInvokedFunctionsPlugin());
+
         return localQueryRunner;
     }
 
diff --git a/presto-tests/src/test/java/com/facebook/presto/tests/TestQueryPlanDeterminism.java b/presto-tests/src/test/java/com/facebook/presto/tests/TestQueryPlanDeterminism.java
index 98fdee5fe7ebf..8fbcc731fec0a 100644
--- a/presto-tests/src/test/java/com/facebook/presto/tests/TestQueryPlanDeterminism.java
+++ b/presto-tests/src/test/java/com/facebook/presto/tests/TestQueryPlanDeterminism.java
@@ -16,6 +16,7 @@
 import com.facebook.presto.Session;
 import com.facebook.presto.common.type.Type;
 import com.facebook.presto.metadata.SessionPropertyManager;
+import com.facebook.presto.scalar.sql.SqlInvokedFunctionsPlugin;
 import com.facebook.presto.spi.ConnectorId;
 import com.facebook.presto.testing.LocalQueryRunner;
 import com.facebook.presto.testing.MaterializedResult;
@@ -74,6 +75,7 @@ protected QueryRunner createQueryRunner()
         sessionPropertyManager.addSystemSessionProperties(TEST_SYSTEM_PROPERTIES);
         sessionPropertyManager.addConnectorSessionProperties(new ConnectorId(TESTING_CATALOG), TEST_CATALOG_PROPERTIES);
 
+        localQueryRunner.installPlugin(new SqlInvokedFunctionsPlugin());
         return localQueryRunner;
     }
 
diff --git a/presto-tests/src/test/java/com/facebook/presto/tests/TestSetWorkerSessionPropertiesExcludingInvalidProperties.java b/presto-tests/src/test/java/com/facebook/presto/tests/TestSetWorkerSessionPropertiesExcludingInvalidProperties.java
index 6732acd3c32e2..a1099d3b60ab7 100644
--- a/presto-tests/src/test/java/com/facebook/presto/tests/TestSetWorkerSessionPropertiesExcludingInvalidProperties.java
+++ b/presto-tests/src/test/java/com/facebook/presto/tests/TestSetWorkerSessionPropertiesExcludingInvalidProperties.java
@@ -51,6 +51,6 @@ public void testSetSessionValidJavaWorkerSessionProperty()
                 "MaterializedResult{rows=[[true]], " +
                         "types=[boolean], " +
                         "setSessionProperties={distinct_aggregation_spill_enabled=false}, " +
-                        "resetSessionProperties=[], updateType=SET SESSION}");
+                        "resetSessionProperties=[], updateType=SET SESSION, clearTransactionId=false}");
     }
 }
diff --git a/presto-tests/src/test/java/com/facebook/presto/tests/TestSetWorkerSessionPropertiesIncludingInvalidProperties.java b/presto-tests/src/test/java/com/facebook/presto/tests/TestSetWorkerSessionPropertiesIncludingInvalidProperties.java
index 2e66c59a1e59d..60d49dc94a40f 100644
--- a/presto-tests/src/test/java/com/facebook/presto/tests/TestSetWorkerSessionPropertiesIncludingInvalidProperties.java
+++ b/presto-tests/src/test/java/com/facebook/presto/tests/TestSetWorkerSessionPropertiesIncludingInvalidProperties.java
@@ -42,7 +42,7 @@ public void testSetSessionValidNativeWorkerSessionProperty()
                 "MaterializedResult{rows=[[true]], " +
                         "types=[boolean], " +
                         "setSessionProperties={native_expression_max_array_size_in_reduce=50000}, " +
-                        "resetSessionProperties=[], updateType=SET SESSION}");
+                        "resetSessionProperties=[], updateType=SET SESSION, clearTransactionId=false}");
     }
 
     @Test
@@ -56,6 +56,6 @@ public void testSetSessionValidJavaWorkerSessionProperty()
                 "MaterializedResult{rows=[[true]], " +
                         "types=[boolean], " +
                         "setSessionProperties={distinct_aggregation_spill_enabled=false}, " +
-                        "resetSessionProperties=[], updateType=SET SESSION}");
+                        "resetSessionProperties=[], updateType=SET SESSION, clearTransactionId=false}");
     }
 }
diff --git a/presto-tests/src/test/java/com/facebook/presto/tests/TestSqlFunctions.java b/presto-tests/src/test/java/com/facebook/presto/tests/TestSqlFunctions.java
index cafb8e59e7fc1..e4ec0ad5e39a8 100644
--- a/presto-tests/src/test/java/com/facebook/presto/tests/TestSqlFunctions.java
+++ b/presto-tests/src/test/java/com/facebook/presto/tests/TestSqlFunctions.java
@@ -18,6 +18,7 @@
 import com.facebook.presto.common.type.TypeSignature;
 import com.facebook.presto.common.type.TypeSignatureParameter;
 import com.facebook.presto.common.type.UserDefinedType;
+import com.facebook.presto.scalar.sql.SqlInvokedFunctionsPlugin;
 import com.facebook.presto.spi.function.Parameter;
 import com.facebook.presto.spi.function.RoutineCharacteristics;
 import com.facebook.presto.spi.function.SqlFunctionId;
@@ -99,7 +100,7 @@ protected QueryRunner createQueryRunner()
             queryRunner.getMetadata().getFunctionAndTypeManager().addUserDefinedType(COUNTRY_ENUM);
 
             queryRunner.execute("CREATE TYPE testing.type.person AS (first_name varchar, last_name varchar, age tinyint, country testing.enum.country)");
-
+            queryRunner.installPlugin(new SqlInvokedFunctionsPlugin());
             return queryRunner;
         }
         catch (Exception e) {
diff --git a/presto-tests/src/test/java/com/facebook/presto/tests/TestVerboseOptimizerInfo.java b/presto-tests/src/test/java/com/facebook/presto/tests/TestVerboseOptimizerInfo.java
index 3cb9f26f1b8fd..add9d0be8cd07 100644
--- a/presto-tests/src/test/java/com/facebook/presto/tests/TestVerboseOptimizerInfo.java
+++ b/presto-tests/src/test/java/com/facebook/presto/tests/TestVerboseOptimizerInfo.java
@@ -16,6 +16,7 @@
 
 import com.facebook.presto.Session;
 import com.facebook.presto.metadata.SessionPropertyManager;
+import com.facebook.presto.scalar.sql.SqlInvokedFunctionsPlugin;
 import com.facebook.presto.spi.ConnectorId;
 import com.facebook.presto.testing.LocalQueryRunner;
 import com.facebook.presto.testing.MaterializedResult;
@@ -71,7 +72,7 @@ public static LocalQueryRunner createLocalQueryRunner()
         SessionPropertyManager sessionPropertyManager = localQueryRunner.getMetadata().getSessionPropertyManager();
         sessionPropertyManager.addSystemSessionProperties(TEST_SYSTEM_PROPERTIES);
         sessionPropertyManager.addConnectorSessionProperties(new ConnectorId(TESTING_CATALOG), TEST_CATALOG_PROPERTIES);
-
+        localQueryRunner.installPlugin(new SqlInvokedFunctionsPlugin());
         return localQueryRunner;
     }
 
diff --git a/presto-thrift-connector/src/main/java/com/facebook/presto/connector/thrift/ThriftMetadata.java b/presto-thrift-connector/src/main/java/com/facebook/presto/connector/thrift/ThriftMetadata.java
index d0291feb4035e..63d499bcdc752 100644
--- a/presto-thrift-connector/src/main/java/com/facebook/presto/connector/thrift/ThriftMetadata.java
+++ b/presto-thrift-connector/src/main/java/com/facebook/presto/connector/thrift/ThriftMetadata.java
@@ -42,7 +42,6 @@
 import com.google.common.cache.CacheBuilder;
 import com.google.common.cache.CacheLoader;
 import com.google.common.cache.LoadingCache;
-import com.google.common.collect.ImmutableList;
 import jakarta.inject.Inject;
 
 import java.util.List;
@@ -110,7 +109,7 @@ public ConnectorTableHandle getTableHandle(ConnectorSession session, SchemaTable
     }
 
     @Override
-    public List getTableLayouts(
+    public ConnectorTableLayoutResult getTableLayoutForConstraint(
             ConnectorSession session,
             ConnectorTableHandle table,
             Constraint constraint,
@@ -122,7 +121,7 @@ public List getTableLayouts(
                 tableHandle.getTableName(),
                 desiredColumns,
                 constraint.getSummary());
-        return ImmutableList.of(new ConnectorTableLayoutResult(new ConnectorTableLayout(layoutHandle), constraint.getSummary()));
+        return new ConnectorTableLayoutResult(new ConnectorTableLayout(layoutHandle), constraint.getSummary());
     }
 
     @Override
diff --git a/presto-tpcds/src/main/java/com/facebook/presto/tpcds/TpcdsMetadata.java b/presto-tpcds/src/main/java/com/facebook/presto/tpcds/TpcdsMetadata.java
index ab57475bb199a..7d9e964e6a541 100644
--- a/presto-tpcds/src/main/java/com/facebook/presto/tpcds/TpcdsMetadata.java
+++ b/presto-tpcds/src/main/java/com/facebook/presto/tpcds/TpcdsMetadata.java
@@ -100,7 +100,7 @@ public ConnectorTableHandle getTableHandle(ConnectorSession session, SchemaTable
     }
 
     @Override
-    public List getTableLayouts(
+    public ConnectorTableLayoutResult getTableLayoutForConstraint(
             ConnectorSession session,
             ConnectorTableHandle table,
             Constraint constraint,
@@ -116,7 +116,7 @@ public List getTableLayouts(
                 Optional.empty(),
                 ImmutableList.of());
 
-        return ImmutableList.of(new ConnectorTableLayoutResult(layout, constraint.getSummary()));
+        return new ConnectorTableLayoutResult(layout, constraint.getSummary());
     }
 
     @Override
@@ -124,8 +124,7 @@ public ConnectorTableLayout getTableLayout(ConnectorSession session, ConnectorTa
     {
         TpcdsTableLayoutHandle layout = (TpcdsTableLayoutHandle) handle;
 
-        return getTableLayouts(session, layout.getTable(), Constraint.alwaysTrue(), Optional.empty())
-                .get(0)
+        return getTableLayoutForConstraint(session, layout.getTable(), Constraint.alwaysTrue(), Optional.empty())
                 .getTableLayout();
     }
 
diff --git a/presto-tpch/src/main/java/com/facebook/presto/tpch/TpchMetadata.java b/presto-tpch/src/main/java/com/facebook/presto/tpch/TpchMetadata.java
index 6e27f7df994b6..255e80cbdafe2 100644
--- a/presto-tpch/src/main/java/com/facebook/presto/tpch/TpchMetadata.java
+++ b/presto-tpch/src/main/java/com/facebook/presto/tpch/TpchMetadata.java
@@ -186,7 +186,7 @@ public ConnectorTableHandle getTableHandleForStatisticsCollection(ConnectorSessi
     }
 
     @Override
-    public List getTableLayouts(
+    public ConnectorTableLayoutResult getTableLayoutForConstraint(
             ConnectorSession session,
             ConnectorTableHandle table,
             Constraint constraint,
@@ -252,7 +252,7 @@ else if (tableHandle.getTableName().equals(TpchTable.LINE_ITEM.getTableName()))
                 Optional.empty(),
                 localProperties);
 
-        return ImmutableList.of(new ConnectorTableLayoutResult(layout, unenforcedConstraint));
+        return new ConnectorTableLayoutResult(layout, unenforcedConstraint);
     }
 
     private Set filterValues(Set nullableValues, TpchColumn column, Constraint constraint)
@@ -269,8 +269,7 @@ public ConnectorTableLayout getTableLayout(ConnectorSession session, ConnectorTa
         TpchTableLayoutHandle layout = (TpchTableLayoutHandle) handle;
 
         // tables in this connector have a single layout
-        return getTableLayouts(session, layout.getTable(), Constraint.alwaysTrue(), Optional.empty())
-                .get(0)
+        return getTableLayoutForConstraint(session, layout.getTable(), Constraint.alwaysTrue(), Optional.empty())
                 .getTableLayout();
     }
 
diff --git a/presto-tpch/src/test/java/com/facebook/presto/tpch/TestTpchMetadata.java b/presto-tpch/src/test/java/com/facebook/presto/tpch/TestTpchMetadata.java
index d6875ae59f612..b7a2737d46029 100644
--- a/presto-tpch/src/test/java/com/facebook/presto/tpch/TestTpchMetadata.java
+++ b/presto-tpch/src/test/java/com/facebook/presto/tpch/TestTpchMetadata.java
@@ -46,7 +46,6 @@
 import static com.facebook.presto.spi.Constraint.alwaysTrue;
 import static com.facebook.presto.tpch.TpchMetadata.getPrestoType;
 import static com.facebook.presto.tpch.util.PredicateUtils.filterOutColumnFromPredicate;
-import static com.google.common.collect.MoreCollectors.onlyElement;
 import static io.airlift.slice.Slices.utf8Slice;
 import static io.airlift.tpch.CustomerColumn.MARKET_SEGMENT;
 import static io.airlift.tpch.CustomerColumn.NAME;
@@ -389,9 +388,7 @@ private static TupleDomain fixedValueTupleDomain(TpchMetadata tpch
 
     private static ConnectorTableLayoutResult getTableOnlyLayout(TpchMetadata tpchMetadata, ConnectorSession session, ConnectorTableHandle tableHandle, Constraint constraint)
     {
-        return tpchMetadata.getTableLayouts(session, tableHandle, constraint, Optional.empty())
-                .stream()
-                .collect(onlyElement());
+        return tpchMetadata.getTableLayoutForConstraint(session, tableHandle, constraint, Optional.empty());
     }
 
     private ColumnStatistics noColumnStatistics()
diff --git a/presto-ui/src/static/oauth2/failure.html b/presto-ui/src/static/oauth2/failure.html
new file mode 100644
index 0000000000000..1f4dc659ad49d
--- /dev/null
+++ b/presto-ui/src/static/oauth2/failure.html
@@ -0,0 +1,60 @@
+
+
+
+
+    
+    
+
+    
+    Failure
+
+    
+
+    
+    
+
+    
+    
+
+    
+    
+
+    
+    
+
+    
+    
+
+    
+    
+
+    
+    
+
+
+
+
+
OAuth2 authentication failed

+
+

+
+
+
+
+
+
diff --git a/presto-ui/src/static/oauth2/logout.html b/presto-ui/src/static/oauth2/logout.html
new file mode 100644
index 0000000000000..821687918d81c
--- /dev/null
+++ b/presto-ui/src/static/oauth2/logout.html
@@ -0,0 +1,51 @@
+
+
+
+
+    
+    
+
+    
+    Logout
+
+    
+
+    
+    
+
+    
+    
+
+    
+    
+
+    
+    
+
+
+
+
+
OAuth2 logout succeeded

+
+
This browser window can be closed

+
+
+
+
+
+
diff --git a/presto-ui/src/static/oauth2/success.html b/presto-ui/src/static/oauth2/success.html
new file mode 100644
index 0000000000000..dac401a1d9e6d
--- /dev/null
+++ b/presto-ui/src/static/oauth2/success.html
@@ -0,0 +1,60 @@
+
+
+
+
+    
+    
+
+    
+    Success
+
+    
+
+    
+    
+
+    
+    
+
+    
+    
+
+    
+    
+
+    
+    
+
+    
+    
+
+    
+    
+
+
+
+
+
OAuth2 authentication succeeded

+
+
This browser window can be closed

+
+
+
+
+
+