Pushdown integer upcasts to source operators in LocalPlanner

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

@@ -599,7 +599,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

@@ -58,15 +58,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(

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;

velox/connectors/hive/HiveDataSource.cpp

@@ -91,14 +91,17 @@ HiveDataSource::HiveDataSource(
     FileHandleFactory* fileHandleFactory,
     folly::Executor* ioExecutor,
     const ConnectorQueryCtx* connectorQueryCtx,
-    const std::shared_ptr<HiveConfig>& hiveConfig)
-    : fileHandleFactory_(fileHandleFactory),
+    const std::shared_ptr<HiveConfig>& hiveConfig,
+    bool pushdownCasts)
+    : assignments_(assignments),
+      fileHandleFactory_(fileHandleFactory),
       ioExecutor_(ioExecutor),
       connectorQueryCtx_(connectorQueryCtx),
       hiveConfig_(hiveConfig),
       pool_(connectorQueryCtx->memoryPool()),
       outputType_(outputType),
-      expressionEvaluator_(connectorQueryCtx->expressionEvaluator()) {
+      expressionEvaluator_(connectorQueryCtx->expressionEvaluator()),
+      pushdownCasts_(pushdownCasts) {
   hiveTableHandle_ =
       std::dynamic_pointer_cast<const HiveTableHandle>(tableHandle);
   VELOX_CHECK_NOT_NULL(
@@ -138,22 +141,44 @@ HiveDataSource::HiveDataSource(
   }
 
   std::vector<std::string> readColumnNames;
-  auto readColumnTypes = outputType_->children();
-  for (const auto& outputName : outputType_->names()) {
-    auto it = assignments.find(outputName);
+  std::vector<TypePtr> readColumnTypes;
+  std::vector<std::string> readColumnNamesWithoutUpcasts;
+  std::vector<TypePtr> readColumnTypesWithoutUpcasts;
+
+  // outputType_ contains the upcast columns if pushdownCasts_ is true.
+  for (int i = 0; i < outputType_->size(); ++i) {
+    auto columnName = outputType_->nameOf(i); // e.g. order_id_21_upcast
+    auto& columnType = outputType_->childAt(i);
+
+    auto originalColumnName = columnName;
+    if (pushdownCasts_ && columnName.ends_with("_upcast")) {
+      originalColumnName =
+          columnName.substr(0, columnName.size() - strlen("_upcast"));
+    }
+
+    // Get the ColumnHandle name. This is the name without aliasing. e.g.
+    // originalColumnName="order_id_21", and columnHandleName="order_id"
+    auto it = assignments_.find(originalColumnName);
     VELOX_CHECK(
-        it != assignments.end(),
+        it != assignments_.end(),
         "ColumnHandle is missing for output column: {}",
-        outputName);
-
+        columnName);
     auto* handle = static_cast<const HiveColumnHandle*>(it->second.get());
-    readColumnNames.push_back(handle->name());
+    auto columnHandleName = handle->name();
+
+    if (!pushdownCasts_ || !columnName.ends_with("_upcast")) {
+      readColumnNamesWithoutUpcasts.push_back(columnHandleName);
+      readColumnTypesWithoutUpcasts.push_back(columnType);
+    }
+    readColumnNames.push_back(columnHandleName);
+    readColumnTypes.push_back(columnType);
+
     for (auto& subfield : handle->requiredSubfields()) {
       VELOX_USER_CHECK_EQ(
           getColumnName(subfield),
           handle->name(),
           "Required subfield does not match column name");
-      subfields_[handle->name()].push_back(&subfield);
+      subfields_[columnHandleName].push_back(&subfield);
     }
     columnPostProcessors_.push_back(handle->postProcessor());
   }
@@ -195,7 +220,7 @@ HiveDataSource::HiveDataSource(
       }
       // Remaining filter may reference columns that are not used otherwise,
       // e.g. are not being projected out and are not used in range filters.
-      // Make sure to add these columns to readerOutputType_.
+      // Make sure to add these columns to readerOutputTypeWithoutUpcasts_.
       readColumnNames.push_back(input->field());
       readColumnTypes.push_back(input->type());
     }
@@ -222,8 +247,12 @@ HiveDataSource::HiveDataSource(
 
   readerOutputType_ =
       ROW(std::move(readColumnNames), std::move(readColumnTypes));
+  // NO upcast columns
+  readerOutputTypeWithoutUpcasts_ =
+      ROW(std::move(readColumnNamesWithoutUpcasts),
+          std::move(readColumnTypesWithoutUpcasts));
   scanSpec_ = makeScanSpec(
-      readerOutputType_,
+      readerOutputTypeWithoutUpcasts_,
       subfields_,
       filters_,
       hiveTableHandle_->dataColumns(),
@@ -249,7 +278,7 @@ std::unique_ptr<SplitReader> HiveDataSource::createSplitReader() {
       &partitionKeys_,
       connectorQueryCtx_,
       hiveConfig_,
-      readerOutputType_,
+      readerOutputTypeWithoutUpcasts_,
       ioStats_,
       fsStats_,
       fileHandleFactory_,
@@ -273,11 +302,12 @@ std::vector<column_index_t> HiveDataSource::setupBucketConversion() {
     if (subfields_.erase(handle->name()) > 0) {
       rebuildScanSpec = true;
     }
-    auto index = readerOutputType_->getChildIdxIfExists(handle->name());
+    auto index =
+        readerOutputTypeWithoutUpcasts_->getChildIdxIfExists(handle->name());
     if (!index.has_value()) {
       if (names.empty()) {
-        names = readerOutputType_->names();
-        types = readerOutputType_->children();
+        names = readerOutputTypeWithoutUpcasts_->names();
+        types = readerOutputTypeWithoutUpcasts_->children();
       }
       index = names.size();
       names.push_back(handle->name());
@@ -288,11 +318,11 @@ std::vector<column_index_t> HiveDataSource::setupBucketConversion() {
     bucketChannels.push_back(*index);
   }
   if (!names.empty()) {
-    readerOutputType_ = ROW(std::move(names), std::move(types));
+    readerOutputTypeWithoutUpcasts_ = ROW(std::move(names), std::move(types));
   }
   if (rebuildScanSpec) {
     auto newScanSpec = makeScanSpec(
-        readerOutputType_,
+        readerOutputTypeWithoutUpcasts_,
         subfields_,
         filters_,
         hiveTableHandle_->dataColumns(),
@@ -314,7 +344,8 @@ void HiveDataSource::setupRowIdColumn() {
   auto* rowId = scanSpec_->childByName(*specialColumns_.rowId);
   VELOX_CHECK_NOT_NULL(rowId);
   auto& rowIdType =
-      readerOutputType_->findChild(*specialColumns_.rowId)->asRow();
+      readerOutputTypeWithoutUpcasts_->findChild(*specialColumns_.rowId)
+          ->asRow();
   auto rowGroupId = split_->getFileName();
   rowId->childByName(rowIdType.nameOf(1))
       ->setConstantValue<StringView>(
@@ -339,8 +370,6 @@ void HiveDataSource::addSplit(std::shared_ptr<ConnectorSplit> split) {
   split_ = std::dynamic_pointer_cast<HiveConnectorSplit>(split);
   VELOX_CHECK_NOT_NULL(split_, "Wrong type of split");
 
-  VLOG(1) << "Adding split " << split_->toString();
-
   if (splitReader_) {
     splitReader_.reset();
   }
@@ -361,7 +390,7 @@ void HiveDataSource::addSplit(std::shared_ptr<ConnectorSplit> split) {
   // so we initialize it beforehand.
   splitReader_->configureReaderOptions(randomSkip_);
   splitReader_->prepareSplit(metadataFilter_, runtimeStats_);
-  readerOutputType_ = splitReader_->readerOutputType();
+  readerOutputTypeWithoutUpcasts_ = splitReader_->readerOutputType();
 }
 
 std::optional<RowVectorPtr> HiveDataSource::next(
@@ -380,14 +409,16 @@ std::optional<RowVectorPtr> HiveDataSource::next(
 
   // Bucket conversion or delta update could add extra column to reader output.
   auto needsExtraColumn = [&] {
-    return output_->asUnchecked<RowVector>()->childrenSize() <
-        readerOutputType_->size();
+    return outputWithoutUpcasts_->asUnchecked<RowVector>()->childrenSize() <
+        readerOutputTypeWithoutUpcasts_->size();
   };
-  if (!output_ || needsExtraColumn()) {
-    output_ = BaseVector::create(readerOutputType_, 0, pool_);
+  if (!outputWithoutUpcasts_ || needsExtraColumn()) {
+    outputWithoutUpcasts_ =
+        BaseVector::create(readerOutputTypeWithoutUpcasts_, 0, pool_);
   }
 
-  const auto rowsScanned = splitReader_->next(size, output_);
+  // Read only the real columns, not the upcast columns.
+  const auto rowsScanned = splitReader_->next(size, outputWithoutUpcasts_);
   completedRows_ += rowsScanned;
   if (rowsScanned == 0) {
     splitReader_->updateRuntimeStats(runtimeStats_);
@@ -396,14 +427,15 @@ std::optional<RowVectorPtr> HiveDataSource::next(
   }
 
   VELOX_CHECK(
-      !output_->mayHaveNulls(), "Top-level row vector cannot have nulls");
-  auto rowsRemaining = output_->size();
+      !outputWithoutUpcasts_->mayHaveNulls(),
+      "Top-level row vector cannot have nulls");
+  auto rowsRemaining = outputWithoutUpcasts_->size();
   if (rowsRemaining == 0) {
     // no rows passed the pushed down filters.
     return getEmptyOutput();
   }
 
-  auto rowVector = std::dynamic_pointer_cast<RowVector>(output_);
+  auto rowVector = std::dynamic_pointer_cast<RowVector>(outputWithoutUpcasts_);
 
   // In case there is a remaining filter that excludes some but not all
   // rows, collect the indices of the passing rows. If there is no filter,
@@ -433,12 +465,49 @@ std::optional<RowVectorPtr> HiveDataSource::next(
   std::vector<VectorPtr> outputColumns;
   outputColumns.reserve(outputType_->size());
   for (int i = 0; i < outputType_->size(); ++i) {
-    auto& child = rowVector->childAt(i);
-    if (remainingIndices) {
-      // Disable dictionary values caching in expression eval so that we
-      // don't need to reallocate the result for every batch.
-      child->disableMemo();
+    std::shared_ptr<BaseVector> child;
+    // find the upcast columns and add them to outputWithoutUpcasts_
+    const auto& columnName = outputType_->nameOf(i);
+    // outputType_ includes the upcast columns,
+    const auto& columnType = outputType_->childAt(i);
+
+    if (columnName.ends_with("_upcast")) {
+      auto originalOutputName =
+          columnName.substr(0, columnName.size() - strlen("_upcast"));
+      auto columnHandleIt = assignments_.find(originalOutputName);
+      VELOX_CHECK(
+          columnHandleIt != assignments_.end(),
+          "Cannot find column handle for upcast column: {} original: {}",
+          columnName,
+          originalOutputName);
+      auto columnHandleName =
+          static_cast<const HiveColumnHandle*>(columnHandleIt->second.get())
+              ->name();
+
+      //  rowVector does not have the upcast columns.
+      auto index = readerOutputTypeWithoutUpcasts_->getChildIdxIfExists(
+          columnHandleName);
+      VELOX_CHECK(index.has_value());
+      auto originalColumn = rowVector->childAt(*index);
+
+      child = BaseVector::create(columnType, originalColumn->size(), pool_);
+      child->copy(originalColumn.get(), 0, 0, originalColumn->size());
+    } else {
+      auto columnHandleIt = assignments_.find(columnName);
+      VELOX_CHECK(
+          columnHandleIt != assignments_.end(),
+          "Cannot find column handle for upcast column: {} original: {}",
+          columnName,
+          columnName);
+      auto columnHandleName =
+          static_cast<const HiveColumnHandle*>(columnHandleIt->second.get())
+              ->name();
+      auto index =
+          readerOutputTypeWithoutUpcasts_->getChildIdxIfExists(columnHandleName);
+      VELOX_CHECK(index.has_value());
+      child = rowVector->childAt(*index);
     }
+
     auto column = exec::wrapChild(rowsRemaining, remainingIndices, child);
     if (columnPostProcessors_[i]) {
       columnPostProcessors_[i](column);
@@ -541,7 +610,8 @@ void HiveDataSource::setFromDataSource(
   runtimeStats_.skippedSplits += source->runtimeStats_.skippedSplits;
   runtimeStats_.processedSplits += source->runtimeStats_.processedSplits;
   runtimeStats_.skippedSplitBytes += source->runtimeStats_.skippedSplitBytes;
-  readerOutputType_ = std::move(source->readerOutputType_);
+  readerOutputTypeWithoutUpcasts_ =
+      std::move(source->readerOutputTypeWithoutUpcasts_);
   source->scanSpec_->moveAdaptationFrom(*scanSpec_);
   scanSpec_ = std::move(source->scanSpec_);
   metadataFilter_ = std::move(source->metadataFilter_);
@@ -600,7 +670,7 @@ std::shared_ptr<wave::WaveDataSource> HiveDataSource::toWaveDataSource() {
     waveDataSource_ = waveDelegateHook_(
         hiveTableHandle_,
         scanSpec_,
-        readerOutputType_,
+        readerOutputTypeWithoutUpcasts_,
         &partitionKeys_,
         fileHandleFactory_,
         ioExecutor_,

velox/connectors/hive/HiveDataSource.h

@@ -41,7 +41,8 @@ class HiveDataSource : public DataSource {
       FileHandleFactory* fileHandleFactory,
       folly::Executor* ioExecutor,
       const ConnectorQueryCtx* connectorQueryCtx,
-      const std::shared_ptr<HiveConfig>& hiveConfig);
+      const std::shared_ptr<HiveConfig>& hiveConfig,
+      bool pushdownCasts = false);
 
   void addSplit(std::shared_ptr<ConnectorSplit> split) override;
 
@@ -100,7 +101,7 @@ class HiveDataSource : public DataSource {
 
  protected:
   virtual std::unique_ptr<SplitReader> createSplitReader();
-
+  const connector::ColumnHandleMap assignments_;
   FileHandleFactory* const fileHandleFactory_;
   folly::Executor* const ioExecutor_;
   const ConnectorQueryCtx* const connectorQueryCtx_;
@@ -110,13 +111,15 @@ class HiveDataSource : public DataSource {
   std::shared_ptr<HiveConnectorSplit> split_;
   HiveTableHandlePtr hiveTableHandle_;
   std::shared_ptr<common::ScanSpec> scanSpec_;
+  VectorPtr outputWithoutUpcasts_;
   VectorPtr output_;
   std::unique_ptr<SplitReader> splitReader_;
-
   // Output type from file reader.  This is different from outputType_ that it
   // contains column names before assignment, and columns that only used in
   // remaining filter.
   RowTypePtr readerOutputType_;
+  // The ColumnHandle name, e.g. order_id
+  RowTypePtr readerOutputTypeWithoutUpcasts_;
 
   // Column handles for the partition key columns keyed on partition key column
   // name.
@@ -151,9 +154,15 @@ class HiveDataSource : public DataSource {
   // object.
   void processColumnHandle(const HiveColumnHandlePtr& handle);
 
-  // The row type for the data source output, not including filter-only columns
+  // The row type for the data source output, including filter-only columns
+  // May be aliased, e.g. (order_id_21, order_id_21_upcast). Does not include
+  // filter only columns
   const RowTypePtr outputType_;
+  // Same as outputType_ but the column names are the ColumnHandle names
+  RowTypePtr outputTypeWithoutUpcasts_;
+
   core::ExpressionEvaluator* const expressionEvaluator_;
+  const bool pushdownCasts_;
 
   // Column handles for the Split info columns keyed on their column names.
   std::unordered_map<std::string, HiveColumnHandlePtr> infoColumns_;

velox/connectors/tests/ConnectorTest.cpp

@@ -30,7 +30,8 @@ class TestConnector : public connector::Connector {
       const RowTypePtr& /* outputType */,
       const ConnectorTableHandlePtr& /* tableHandle */,
       const connector::ColumnHandleMap& /* columnHandles */,
-      connector::ConnectorQueryCtx* connectorQueryCtx) override {
+      connector::ConnectorQueryCtx* connectorQueryCtx,
+      bool pushdownCasts) override {
     VELOX_NYI();
   }