feat: Add probe/build geometries and radius as members of SpatialJoinNode

presto-main-base/src/main/java/com/facebook/presto/sql/planner/iterative/rule/ExtractSpatialJoins.java

@@ -460,6 +460,8 @@ else if (firstSide == VariableSide.Right && secondSide == VariableSide.Left) {
         // with a projection that adds the argument as a variable.
         Optional<VariableReferenceExpression> newFirstVariable = newVariable(context, firstArgument);
         Optional<VariableReferenceExpression> newSecondVariable = newVariable(context, secondArgument);
+        VariableReferenceExpression leftGeometryVariable;
+        VariableReferenceExpression rightGeometryVariable;
 
         PlanNode leftNode = joinNode.getLeft();
         PlanNode rightNode = joinNode.getRight();
@@ -470,10 +472,16 @@ else if (firstSide == VariableSide.Right && secondSide == VariableSide.Left) {
         if (firstArgumentOnLeft) {
             newLeftNode = newFirstVariable.map(variable -> addProjection(context, leftNode, variable, firstArgument)).orElse(leftNode);
             newRightNode = newSecondVariable.map(variable -> addProjection(context, rightNode, variable, secondArgument)).orElse(rightNode);
+            // If new variables are empty, argument is VariableReferenceExpression
+            leftGeometryVariable = newFirstVariable.orElseGet(() -> (VariableReferenceExpression) firstArgument);
+            rightGeometryVariable = newSecondVariable.orElseGet(() -> (VariableReferenceExpression) secondArgument);
         }
         else {
             newLeftNode = newSecondVariable.map(variable -> addProjection(context, leftNode, variable, secondArgument)).orElse(leftNode);
             newRightNode = newFirstVariable.map(variable -> addProjection(context, rightNode, variable, firstArgument)).orElse(rightNode);
+            // If new variables are empty, argument is VariableReferenceExpression
+            leftGeometryVariable = newSecondVariable.orElseGet(() -> (VariableReferenceExpression) secondArgument);
+            rightGeometryVariable = newFirstVariable.orElseGet(() -> (VariableReferenceExpression) firstArgument);
         }
 
         RowExpression newFirstArgument = mapToExpression(newFirstVariable, firstArgument);
@@ -512,6 +520,9 @@ else if (firstSide == VariableSide.Right && secondSide == VariableSide.Left) {
                 newLeftNode,
                 newRightNode,
                 outputVariables,
+                leftGeometryVariable,
+                rightGeometryVariable,
+                radius,
                 newFilter,
                 leftPartitionVariable,
                 rightPartitionVariable,

presto-main-base/src/main/java/com/facebook/presto/sql/planner/iterative/rule/RowExpressionRewriteRuleSet.java

@@ -218,6 +218,9 @@ public Result apply(SpatialJoinNode spatialJoinNode, Captures captures, Context
                     spatialJoinNode.getLeft(),
                     spatialJoinNode.getRight(),
                     spatialJoinNode.getOutputVariables(),
+                    spatialJoinNode.getProbeGeometryVariable(),
+                    spatialJoinNode.getBuildGeometryVariable(),
+                    spatialJoinNode.getRadiusVariable(),
                     rewritten,
                     spatialJoinNode.getLeftPartitionVariable(),
                     spatialJoinNode.getRightPartitionVariable(),

presto-main-base/src/main/java/com/facebook/presto/sql/planner/optimizations/PredicatePushDown.java

@@ -914,6 +914,9 @@ public PlanNode visitSpatialJoin(SpatialJoinNode node, RewriteContext<RowExpress
                         node.getLeft(),
                         node.getRight(),
                         node.getOutputVariables(),
+                        node.getProbeGeometryVariable(),
+                        node.getBuildGeometryVariable(),
+                        node.getRadiusVariable(),
                         node.getFilter(),
                         node.getLeftPartitionVariable(),
                         node.getRightPartitionVariable(),
@@ -988,6 +991,9 @@ public PlanNode visitSpatialJoin(SpatialJoinNode node, RewriteContext<RowExpress
                         leftSource,
                         rightSource,
                         node.getOutputVariables(),
+                        node.getProbeGeometryVariable(),
+                        node.getBuildGeometryVariable(),
+                        node.getRadiusVariable(),
                         newJoinPredicate,
                         node.getLeftPartitionVariable(),
                         node.getRightPartitionVariable(),

presto-main-base/src/main/java/com/facebook/presto/sql/planner/optimizations/PruneUnreferencedOutputs.java

@@ -320,7 +320,21 @@ public PlanNode visitSpatialJoin(SpatialJoinNode node, RewriteContext<Set<Variab
 
             planChanged = outputVariables.size() != node.getOutputVariables().size();
 
-            return new SpatialJoinNode(node.getSourceLocation(), node.getId(), node.getStatsEquivalentPlanNode(), node.getType(), left, right, outputVariables, node.getFilter(), node.getLeftPartitionVariable(), node.getRightPartitionVariable(), node.getKdbTree());
+            return new SpatialJoinNode(
+                    node.getSourceLocation(),
+                    node.getId(),
+                    node.getStatsEquivalentPlanNode(),
+                    node.getType(),
+                    left,
+                    right,
+                    outputVariables,
+                    node.getProbeGeometryVariable(),
+                    node.getBuildGeometryVariable(),
+                    node.getRadiusVariable(),
+                    node.getFilter(),
+                    node.getLeftPartitionVariable(),
+                    node.getRightPartitionVariable(),
+                    node.getKdbTree());
         }
 
         @Override

presto-main-base/src/main/java/com/facebook/presto/sql/planner/optimizations/UnaliasSymbolReferences.java

@@ -653,7 +653,20 @@ public PlanNode visitSpatialJoin(SpatialJoinNode node, RewriteContext<Void> cont
             PlanNode left = context.rewrite(node.getLeft());
             PlanNode right = context.rewrite(node.getRight());
 
-            return new SpatialJoinNode(node.getSourceLocation(), node.getId(), node.getType(), left, right, canonicalizeAndDistinct(node.getOutputVariables()), canonicalize(node.getFilter()), canonicalize(node.getLeftPartitionVariable()), canonicalize(node.getRightPartitionVariable()), node.getKdbTree());
+            return new SpatialJoinNode(
+                    node.getSourceLocation(),
+                    node.getId(),
+                    node.getType(),
+                    left,
+                    right,
+                    canonicalizeAndDistinct(node.getOutputVariables()),
+                    canonicalize(node.getProbeGeometryVariable()),
+                    canonicalize(node.getBuildGeometryVariable()),
+                    canonicalize(node.getRadiusVariable()),
+                    canonicalize(node.getFilter()),
+                    canonicalize(node.getLeftPartitionVariable()),
+                    canonicalize(node.getRightPartitionVariable()),
+                    node.getKdbTree());
         }
 
         @Override

presto-native-execution/presto_cpp/main/types/PrestoToVeloxQueryPlan.cpp

@@ -1284,11 +1284,18 @@ core::PlanNodePtr VeloxQueryPlanConverterBase::toVeloxQueryPlan(
     const std::shared_ptr<protocol::TableWriteInfo>& tableWriteInfo,
     const protocol::TaskId& taskId) {
   auto joinType = toJoinType(node->type);
+  std::optional<core::FieldAccessTypedExprPtr> radiusVariable = std::nullopt;
+  if (node->radiusVariable) {
+    radiusVariable = exprConverter_.toVeloxExpr(*node->radiusVariable);
+  }
 
   return std::make_shared<core::SpatialJoinNode>(
       node->id,
       joinType,
       exprConverter_.toVeloxExpr(node->filter),
+      exprConverter_.toVeloxExpr(node->probeGeometryVariable),
+      exprConverter_.toVeloxExpr(node->buildGeometryVariable),
+      radiusVariable,
       toVeloxQueryPlan(node->left, tableWriteInfo, taskId),
       toVeloxQueryPlan(node->right, tableWriteInfo, taskId),
       toRowType(node->outputVariables, typeParser_));

presto-native-execution/presto_cpp/presto_protocol/connector/hive/presto_protocol_hive.cpp

@@ -370,9 +370,10 @@ namespace facebook::presto::protocol::hive {
 
 // NOLINTNEXTLINE: cppcoreguidelines-avoid-c-arrays
 static const std::pair<BucketFunctionType, json>
-    BucketFunctionType_enum_table[] = { // NOLINT: cert-err58-cpp
-        {BucketFunctionType::HIVE_COMPATIBLE, "HIVE_COMPATIBLE"},
-        {BucketFunctionType::PRESTO_NATIVE, "PRESTO_NATIVE"}};
+    BucketFunctionType_enum_table[] =
+        { // NOLINT: cert-err58-cpp
+            {BucketFunctionType::HIVE_COMPATIBLE, "HIVE_COMPATIBLE"},
+            {BucketFunctionType::PRESTO_NATIVE, "PRESTO_NATIVE"}};
 void to_json(json& j, const BucketFunctionType& e) {
   static_assert(
       std::is_enum<BucketFunctionType>::value,
@@ -598,12 +599,13 @@ namespace facebook::presto::protocol::hive {
 
 // NOLINTNEXTLINE: cppcoreguidelines-avoid-c-arrays
 static const std::pair<HiveCompressionCodec, json>
-    HiveCompressionCodec_enum_table[] = { // NOLINT: cert-err58-cpp
-        {HiveCompressionCodec::NONE, "NONE"},
-        {HiveCompressionCodec::SNAPPY, "SNAPPY"},
-        {HiveCompressionCodec::GZIP, "GZIP"},
-        {HiveCompressionCodec::LZ4, "LZ4"},
-        {HiveCompressionCodec::ZSTD, "ZSTD"}};
+    HiveCompressionCodec_enum_table[] =
+        { // NOLINT: cert-err58-cpp
+            {HiveCompressionCodec::NONE, "NONE"},
+            {HiveCompressionCodec::SNAPPY, "SNAPPY"},
+            {HiveCompressionCodec::GZIP, "GZIP"},
+            {HiveCompressionCodec::LZ4, "LZ4"},
+            {HiveCompressionCodec::ZSTD, "ZSTD"}};
 void to_json(json& j, const HiveCompressionCodec& e) {
   static_assert(
       std::is_enum<HiveCompressionCodec>::value,

presto-native-execution/presto_cpp/presto_protocol/connector/iceberg/presto_protocol_iceberg.cpp

@@ -25,11 +25,12 @@ namespace facebook::presto::protocol::iceberg {
 
 // NOLINTNEXTLINE: cppcoreguidelines-avoid-c-arrays
 static const std::pair<ChangelogOperation, json>
-    ChangelogOperation_enum_table[] = { // NOLINT: cert-err58-cpp
-        {ChangelogOperation::INSERT, "INSERT"},
-        {ChangelogOperation::DELETE, "DELETE"},
-        {ChangelogOperation::UPDATE_BEFORE, "UPDATE_BEFORE"},
-        {ChangelogOperation::UPDATE_AFTER, "UPDATE_AFTER"}};
+    ChangelogOperation_enum_table[] =
+        { // NOLINT: cert-err58-cpp
+            {ChangelogOperation::INSERT, "INSERT"},
+            {ChangelogOperation::DELETE, "DELETE"},
+            {ChangelogOperation::UPDATE_BEFORE, "UPDATE_BEFORE"},
+            {ChangelogOperation::UPDATE_AFTER, "UPDATE_AFTER"}};
 void to_json(json& j, const ChangelogOperation& e) {
   static_assert(
       std::is_enum<ChangelogOperation>::value,
@@ -508,14 +509,15 @@ namespace facebook::presto::protocol::iceberg {
 
 // NOLINTNEXTLINE: cppcoreguidelines-avoid-c-arrays
 static const std::pair<PartitionTransformType, json>
-    PartitionTransformType_enum_table[] = { // NOLINT: cert-err58-cpp
-        {PartitionTransformType::IDENTITY, "IDENTITY"},
-        {PartitionTransformType::YEAR, "YEAR"},
-        {PartitionTransformType::MONTH, "MONTH"},
-        {PartitionTransformType::DAY, "DAY"},
-        {PartitionTransformType::HOUR, "HOUR"},
-        {PartitionTransformType::BUCKET, "BUCKET"},
-        {PartitionTransformType::TRUNCATE, "TRUNCATE"}};
+    PartitionTransformType_enum_table[] =
+        { // NOLINT: cert-err58-cpp
+            {PartitionTransformType::IDENTITY, "IDENTITY"},
+            {PartitionTransformType::YEAR, "YEAR"},
+            {PartitionTransformType::MONTH, "MONTH"},
+            {PartitionTransformType::DAY, "DAY"},
+            {PartitionTransformType::HOUR, "HOUR"},
+            {PartitionTransformType::BUCKET, "BUCKET"},
+            {PartitionTransformType::TRUNCATE, "TRUNCATE"}};
 void to_json(json& j, const PartitionTransformType& e) {
   static_assert(
       std::is_enum<PartitionTransformType>::value,

presto-native-execution/presto_cpp/presto_protocol/core/presto_protocol_core.cpp

@@ -36,10 +36,11 @@ namespace facebook::presto::protocol {
 
 // NOLINTNEXTLINE: cppcoreguidelines-avoid-c-arrays
 static const std::pair<NodeSelectionStrategy, json>
-    NodeSelectionStrategy_enum_table[] = { // NOLINT: cert-err58-cpp
-        {NodeSelectionStrategy::HARD_AFFINITY, "HARD_AFFINITY"},
-        {NodeSelectionStrategy::SOFT_AFFINITY, "SOFT_AFFINITY"},
-        {NodeSelectionStrategy::NO_PREFERENCE, "NO_PREFERENCE"}};
+    NodeSelectionStrategy_enum_table[] =
+        { // NOLINT: cert-err58-cpp
+            {NodeSelectionStrategy::HARD_AFFINITY, "HARD_AFFINITY"},
+            {NodeSelectionStrategy::SOFT_AFFINITY, "SOFT_AFFINITY"},
+            {NodeSelectionStrategy::NO_PREFERENCE, "NO_PREFERENCE"}};
 void to_json(json& j, const NodeSelectionStrategy& e) {
   static_assert(
       std::is_enum<NodeSelectionStrategy>::value,
@@ -558,11 +559,12 @@ namespace facebook::presto::protocol {
 
 // NOLINTNEXTLINE: cppcoreguidelines-avoid-c-arrays
 static const std::pair<AggregationNodeStep, json>
-    AggregationNodeStep_enum_table[] = { // NOLINT: cert-err58-cpp
-        {AggregationNodeStep::PARTIAL, "PARTIAL"},
-        {AggregationNodeStep::FINAL, "FINAL"},
-        {AggregationNodeStep::INTERMEDIATE, "INTERMEDIATE"},
-        {AggregationNodeStep::SINGLE, "SINGLE"}};
+    AggregationNodeStep_enum_table[] =
+        { // NOLINT: cert-err58-cpp
+            {AggregationNodeStep::PARTIAL, "PARTIAL"},
+            {AggregationNodeStep::FINAL, "FINAL"},
+            {AggregationNodeStep::INTERMEDIATE, "INTERMEDIATE"},
+            {AggregationNodeStep::SINGLE, "SINGLE"}};
 void to_json(json& j, const AggregationNodeStep& e) {
   static_assert(
       std::is_enum<AggregationNodeStep>::value,
@@ -2808,10 +2810,11 @@ namespace facebook::presto::protocol {
 
 // NOLINTNEXTLINE: cppcoreguidelines-avoid-c-arrays
 static const std::pair<BuiltInFunctionKind, json>
-    BuiltInFunctionKind_enum_table[] = { // NOLINT: cert-err58-cpp
-        {BuiltInFunctionKind::ENGINE, "ENGINE"},
-        {BuiltInFunctionKind::PLUGIN, "PLUGIN"},
-        {BuiltInFunctionKind::WORKER, "WORKER"}};
+    BuiltInFunctionKind_enum_table[] =
+        { // NOLINT: cert-err58-cpp
+            {BuiltInFunctionKind::ENGINE, "ENGINE"},
+            {BuiltInFunctionKind::PLUGIN, "PLUGIN"},
+            {BuiltInFunctionKind::WORKER, "WORKER"}};
 void to_json(json& j, const BuiltInFunctionKind& e) {
   static_assert(
       std::is_enum<BuiltInFunctionKind>::value,
@@ -6144,9 +6147,10 @@ namespace facebook::presto::protocol {
 
 // NOLINTNEXTLINE: cppcoreguidelines-avoid-c-arrays
 static const std::pair<JoinDistributionType, json>
-    JoinDistributionType_enum_table[] = { // NOLINT: cert-err58-cpp
-        {JoinDistributionType::PARTITIONED, "PARTITIONED"},
-        {JoinDistributionType::REPLICATED, "REPLICATED"}};
+    JoinDistributionType_enum_table[] =
+        { // NOLINT: cert-err58-cpp
+            {JoinDistributionType::PARTITIONED, "PARTITIONED"},
+            {JoinDistributionType::REPLICATED, "REPLICATED"}};
 void to_json(json& j, const JoinDistributionType& e) {
   static_assert(
       std::is_enum<JoinDistributionType>::value,
@@ -8211,14 +8215,17 @@ namespace facebook::presto::protocol {
 
 // NOLINTNEXTLINE: cppcoreguidelines-avoid-c-arrays
 static const std::pair<StageExecutionStrategy, json>
-    StageExecutionStrategy_enum_table[] = { // NOLINT: cert-err58-cpp
-        {StageExecutionStrategy::UNGROUPED_EXECUTION, "UNGROUPED_EXECUTION"},
-        {StageExecutionStrategy::FIXED_LIFESPAN_SCHEDULE_GROUPED_EXECUTION,
-         "FIXED_LIFESPAN_SCHEDULE_GROUPED_EXECUTION"},
-        {StageExecutionStrategy::DYNAMIC_LIFESPAN_SCHEDULE_GROUPED_EXECUTION,
-         "DYNAMIC_LIFESPAN_SCHEDULE_GROUPED_EXECUTION"},
-        {StageExecutionStrategy::RECOVERABLE_GROUPED_EXECUTION,
-         "RECOVERABLE_GROUPED_EXECUTION"}};
+    StageExecutionStrategy_enum_table[] =
+        { // NOLINT: cert-err58-cpp
+            {StageExecutionStrategy::UNGROUPED_EXECUTION,
+             "UNGROUPED_EXECUTION"},
+            {StageExecutionStrategy::FIXED_LIFESPAN_SCHEDULE_GROUPED_EXECUTION,
+             "FIXED_LIFESPAN_SCHEDULE_GROUPED_EXECUTION"},
+            {StageExecutionStrategy::
+                 DYNAMIC_LIFESPAN_SCHEDULE_GROUPED_EXECUTION,
+             "DYNAMIC_LIFESPAN_SCHEDULE_GROUPED_EXECUTION"},
+            {StageExecutionStrategy::RECOVERABLE_GROUPED_EXECUTION,
+             "RECOVERABLE_GROUPED_EXECUTION"}};
 void to_json(json& j, const StageExecutionStrategy& e) {
   static_assert(
       std::is_enum<StageExecutionStrategy>::value,
@@ -9606,6 +9613,27 @@ void to_json(json& j, const SpatialJoinNode& p) {
       "SpatialJoinNode",
       "List<VariableReferenceExpression>",
       "outputVariables");
+  to_json_key(
+      j,
+      "probeGeometryVariable",
+      p.probeGeometryVariable,
+      "SpatialJoinNode",
+      "VariableReferenceExpression",
+      "probeGeometryVariable");
+  to_json_key(
+      j,
+      "buildGeometryVariable",
+      p.buildGeometryVariable,
+      "SpatialJoinNode",
+      "VariableReferenceExpression",
+      "buildGeometryVariable");
+  to_json_key(
+      j,
+      "radiusVariable",
+      p.radiusVariable,
+      "SpatialJoinNode",
+      "VariableReferenceExpression",
+      "radiusVariable");
   to_json_key(
       j, "filter", p.filter, "SpatialJoinNode", "RowExpression", "filter");
   to_json_key(
@@ -9639,6 +9667,27 @@ void from_json(const json& j, SpatialJoinNode& p) {
       "SpatialJoinNode",
       "List<VariableReferenceExpression>",
       "outputVariables");
+  from_json_key(
+      j,
+      "probeGeometryVariable",
+      p.probeGeometryVariable,
+      "SpatialJoinNode",
+      "VariableReferenceExpression",
+      "probeGeometryVariable");
+  from_json_key(
+      j,
+      "buildGeometryVariable",
+      p.buildGeometryVariable,
+      "SpatialJoinNode",
+      "VariableReferenceExpression",
+      "buildGeometryVariable");
+  from_json_key(
+      j,
+      "radiusVariable",
+      p.radiusVariable,
+      "SpatialJoinNode",
+      "VariableReferenceExpression",
+      "radiusVariable");
   from_json_key(
       j, "filter", p.filter, "SpatialJoinNode", "RowExpression", "filter");
   from_json_key(
@@ -9888,12 +9937,13 @@ namespace facebook::presto::protocol {
 
 // NOLINTNEXTLINE: cppcoreguidelines-avoid-c-arrays
 static const std::pair<SystemPartitionFunction, json>
-    SystemPartitionFunction_enum_table[] = { // NOLINT: cert-err58-cpp
-        {SystemPartitionFunction::SINGLE, "SINGLE"},
-        {SystemPartitionFunction::HASH, "HASH"},
-        {SystemPartitionFunction::ROUND_ROBIN, "ROUND_ROBIN"},
-        {SystemPartitionFunction::BROADCAST, "BROADCAST"},
-        {SystemPartitionFunction::UNKNOWN, "UNKNOWN"}};
+    SystemPartitionFunction_enum_table[] =
+        { // NOLINT: cert-err58-cpp
+            {SystemPartitionFunction::SINGLE, "SINGLE"},
+            {SystemPartitionFunction::HASH, "HASH"},
+            {SystemPartitionFunction::ROUND_ROBIN, "ROUND_ROBIN"},
+            {SystemPartitionFunction::BROADCAST, "BROADCAST"},
+            {SystemPartitionFunction::UNKNOWN, "UNKNOWN"}};
 void to_json(json& j, const SystemPartitionFunction& e) {
   static_assert(
       std::is_enum<SystemPartitionFunction>::value,
@@ -9930,13 +9980,14 @@ namespace facebook::presto::protocol {
 
 // NOLINTNEXTLINE: cppcoreguidelines-avoid-c-arrays
 static const std::pair<SystemPartitioning, json>
-    SystemPartitioning_enum_table[] = { // NOLINT: cert-err58-cpp
-        {SystemPartitioning::SINGLE, "SINGLE"},
-        {SystemPartitioning::FIXED, "FIXED"},
-        {SystemPartitioning::SOURCE, "SOURCE"},
-        {SystemPartitioning::SCALED, "SCALED"},
-        {SystemPartitioning::COORDINATOR_ONLY, "COORDINATOR_ONLY"},
-        {SystemPartitioning::ARBITRARY, "ARBITRARY"}};
+    SystemPartitioning_enum_table[] =
+        { // NOLINT: cert-err58-cpp
+            {SystemPartitioning::SINGLE, "SINGLE"},
+            {SystemPartitioning::FIXED, "FIXED"},
+            {SystemPartitioning::SOURCE, "SOURCE"},
+            {SystemPartitioning::SCALED, "SCALED"},
+            {SystemPartitioning::COORDINATOR_ONLY, "COORDINATOR_ONLY"},
+            {SystemPartitioning::ARBITRARY, "ARBITRARY"}};
 void to_json(json& j, const SystemPartitioning& e) {
   static_assert(
       std::is_enum<SystemPartitioning>::value,