Pushdown integer upcasts to source operators in LocalPlanner

Changed pushdown_integer_upcasts_to_scan to pushdown_integer_upcasts_to_source

PUshdown integer upcasts to source operators

Only when there are join nodes or result output nodes in the fragment
that an integer upcast is needed. In this case, pushdown the upcasts to
source operators like TableScanNode or ExchangeNode.

Author: yingsu00

velox/connectors/Connector.h

@@ -604,7 +604,8 @@ class Connector {
       const RowTypePtr& outputType,
       const ConnectorTableHandlePtr& tableHandle,
       const connector::ColumnHandleMap& columnHandles,
-      ConnectorQueryCtx* connectorQueryCtx) = 0;
+      ConnectorQueryCtx* connectorQueryCtx,
+      bool pushdownCasts = false) = 0;
 
   /// Returns true if addSplit of DataSource can use 'dataSource' from
   /// ConnectorSplit in addSplit(). If so, TableScan can preload splits

velox/connectors/fuzzer/FuzzerConnector.h

@@ -111,7 +111,8 @@ class FuzzerConnector final : public Connector {
       const RowTypePtr& outputType,
       const ConnectorTableHandlePtr& tableHandle,
       const connector::ColumnHandleMap& /*columnHandles*/,
-      ConnectorQueryCtx* connectorQueryCtx) override final {
+      ConnectorQueryCtx* connectorQueryCtx,
+      bool pushdownCasts = false) override final {
     return std::make_unique<FuzzerDataSource>(
         outputType, tableHandle, connectorQueryCtx->memoryPool());
   }

velox/connectors/hive/HiveConnector.cpp

@@ -57,15 +57,17 @@ std::unique_ptr<DataSource> HiveConnector::createDataSource(
     const RowTypePtr& outputType,
     const ConnectorTableHandlePtr& tableHandle,
     const std::unordered_map<std::string, ColumnHandlePtr>& columnHandles,
-    ConnectorQueryCtx* connectorQueryCtx) {
+    ConnectorQueryCtx* connectorQueryCtx,
+    bool pushdownCasts) {
   return std::make_unique<HiveDataSource>(
       outputType,
       tableHandle,
       columnHandles,
       &fileHandleFactory_,
       ioExecutor_,
       connectorQueryCtx,
-      hiveConfig_);
+      hiveConfig_,
+      pushdownCasts);
 }
 
 std::unique_ptr<DataSink> HiveConnector::createDataSink(
@@ -192,4 +194,55 @@ void HivePartitionFunctionSpec::registerSerDe() {
       "HivePartitionFunctionSpec", HivePartitionFunctionSpec::deserialize);
 }
 
+std::shared_ptr<core::PartitionFunctionSpec>
+HivePartitionFunctionSpec::rewriteInputType(
+    const RowTypePtr& oldInputType,
+    const RowTypePtr& newInputType) const {
+  // If the layout is identical, no change.
+  if (*oldInputType == *newInputType) {
+    return nullptr;
+  }
+
+  std::vector<column_index_t> newChannels;
+  newChannels.reserve(channels_.size());
+  bool changed = false;
+
+  for (auto channel : channels_) {
+    VELOX_CHECK_LT(
+        channel,
+        oldInputType->size(),
+        "Hive bucket channel index {} is out of bounds for old input type {}",
+        channel,
+        oldInputType->toString());
+
+    const auto& name = oldInputType->nameOf(channel);
+    auto newIndexOpt = newInputType->getChildIdxIfExists(name);
+    VELOX_USER_CHECK(
+        newIndexOpt.has_value(),
+        "Failed to find Hive bucket column '{}' in LocalPartition new input type {}",
+        name,
+        newInputType->toString());
+
+    auto newIndex = static_cast<column_index_t>(*newIndexOpt);
+    newChannels.push_back(newIndex);
+    if (newIndex != channel) {
+      changed = true;
+    }
+  }
+
+  if (!changed) {
+    return nullptr;
+  }
+
+  // If we don't have a bucketToPartition map yet, keep using the "late binding"
+  // constructor. Otherwise preserve the bucketToPartition_.
+  if (bucketToPartition_.empty()) {
+    return std::make_shared<HivePartitionFunctionSpec>(
+        numBuckets_, std::move(newChannels), constValues_);
+  }
+
+  return std::make_shared<HivePartitionFunctionSpec>(
+      numBuckets_, bucketToPartition_, std::move(newChannels), constValues_);
+}
+
 } // namespace facebook::velox::connector::hive

velox/connectors/hive/HiveConnector.h

@@ -42,7 +42,8 @@ class HiveConnector : public Connector {
       const RowTypePtr& outputType,
       const ConnectorTableHandlePtr& tableHandle,
       const connector::ColumnHandleMap& columnHandles,
-      ConnectorQueryCtx* connectorQueryCtx) override;
+      ConnectorQueryCtx* connectorQueryCtx,
+      bool pushdownCasts = false) override;
 
   bool supportsSplitPreload() const override {
     return true;
@@ -138,6 +139,10 @@ class HivePartitionFunctionSpec : public core::PartitionFunctionSpec {
 
   static void registerSerDe();
 
+  std::shared_ptr<core::PartitionFunctionSpec> rewriteInputType(
+      const RowTypePtr& oldInputType,
+      const RowTypePtr& newInputType) const override;
+
  private:
   const int numBuckets_;
   const std::vector<int> bucketToPartition_;