fix

Author: ivanmorozov333

ydb/core/formats/arrow/program/abstract.cpp

@@ -27,6 +27,10 @@ TConclusionStatus IResourceProcessor::Execute(const std::shared_ptr<TAccessorsCo
     return DoExecute(resources);
 }
 
+bool IResourceProcessor::DoHasExecutionData(const ui32 columnId, const std::shared_ptr<TAccessorsCollection>& resources) const {
+    return resources->HasColumn(columnId);
+}
+
 NJson::TJsonValue TResourceProcessorStep::DebugJson() const {
     NJson::TJsonValue result = NJson::JSON_MAP;
     if (ColumnsToFetch.size()) {

ydb/core/formats/arrow/program/abstract.h

@@ -50,6 +50,10 @@ class IColumnResolver {
     virtual ~IColumnResolver() = default;
     virtual TString GetColumnName(ui32 id, bool required = true) const = 0;
     virtual std::optional<ui32> GetColumnIdOptional(const TString& name) const = 0;
+    bool HasColumn(const ui32 id) const {
+        return !!GetColumnName(id, false);
+    }
+
     ui32 GetColumnIdVerified(const char* name) const {
         auto result = GetColumnIdOptional(name);
         AFL_VERIFY(!!result);
@@ -174,12 +178,23 @@ class IResourceProcessor {
     virtual NJson::TJsonValue DoDebugJson() const {
         return NJson::JSON_MAP;
     }
+    virtual bool DoHasExecutionData(const ui32 columnId, const std::shared_ptr<TAccessorsCollection>& resources) const;
 
 public:
+    virtual bool IsAggregation() const = 0;
+
     virtual ~IResourceProcessor() = default;
 
+    virtual TString GetKernelClassNameDef(const TString& defaultValue) const {
+        return defaultValue;
+    }
+
     NJson::TJsonValue DebugJson() const;
 
+    bool HasExecutionData(const ui32 columnId, const std::shared_ptr<TAccessorsCollection>& resources) const {
+        return DoHasExecutionData(columnId, resources);
+    }
+
     ui32 GetOutputColumnIdOnce() const {
         AFL_VERIFY(Output.size() == 1)("size", Output.size());
         return Output.front().GetColumnId();
@@ -202,15 +217,17 @@ class IResourceProcessor {
 class TResourceProcessorStep {
 private:
     YDB_READONLY_DEF(std::vector<TColumnChainInfo>, ColumnsToFetch);
+    YDB_READONLY_DEF(std::vector<TColumnChainInfo>, OriginalColumnsToUse);
     YDB_READONLY_DEF(std::shared_ptr<IResourceProcessor>, Processor);
     YDB_READONLY_DEF(std::vector<TColumnChainInfo>, ColumnsToDrop);
 
 public:
     NJson::TJsonValue DebugJson() const;
 
-    TResourceProcessorStep(
-        std::vector<TColumnChainInfo>&& toFetch, std::shared_ptr<IResourceProcessor>&& processor, std::vector<TColumnChainInfo>&& toDrop)
+    TResourceProcessorStep(std::vector<TColumnChainInfo>&& toFetch, std::vector<TColumnChainInfo>&& originalToUse,
+        std::shared_ptr<IResourceProcessor>&& processor, std::vector<TColumnChainInfo>&& toDrop)
         : ColumnsToFetch(std::move(toFetch))
+        , OriginalColumnsToUse(std::move(originalToUse))
         , Processor(std::move(processor))
         , ColumnsToDrop(std::move(toDrop)) {
         AFL_VERIFY(Processor);

ydb/core/formats/arrow/program/aggr_keys.h

@@ -44,6 +44,10 @@ class TAggregateFunction: public TInternalFunction {
     }
 
 public:
+    virtual bool IsAggregation() const override {
+        return true;
+    }
+
     TAggregateFunction(const EAggregate aggregationType, const std::shared_ptr<arrow::compute::FunctionOptions>& functionOptions = nullptr)
         : TBase(functionOptions, true)
         , AggregationType(aggregationType) {
@@ -153,6 +157,9 @@ class TWithKeysAggregationProcessor: public IResourceProcessor {
         , AggregationKeys(std::move(aggregationKeys))
         , Aggregations(std::move(aggregations)) {
     }
+    virtual bool IsAggregation() const override {
+        return true;
+    }
 
 public:
     static const char* GetHouseGroupByName() {

ydb/core/formats/arrow/program/assign_const.h

@@ -10,6 +10,10 @@ class TConstProcessor: public IResourceProcessor {
 
     virtual TConclusionStatus DoExecute(const std::shared_ptr<TAccessorsCollection>& resources) const override;
 
+    virtual bool IsAggregation() const override {
+        return false;
+    }
+
 public:
     TConstProcessor(const std::shared_ptr<arrow::Scalar>& scalar, const ui32 columnId)
         : TBase(std::vector<TColumnChainInfo>(), std::vector<TColumnChainInfo>({ TColumnChainInfo(columnId) }), EProcessorType::Const)

ydb/core/formats/arrow/program/assign_internal.h

@@ -23,6 +23,14 @@ class TCalculationProcessor: public IResourceProcessor {
         , Function(function) {
     }
 
+    virtual TString GetKernelClassNameDef(const TString& defaultValue) const override {
+        return KernelLogic ? KernelLogic->GetClassName() : defaultValue;
+    }
+
+    virtual bool IsAggregation() const override {
+        return Function->IsAggregation();
+    }
+
 public:
     static TConclusion<std::shared_ptr<TCalculationProcessor>> Build(std::vector<TColumnChainInfo>&& input, const TColumnChainInfo& output, 
         const std::shared_ptr<IStepFunction>& function, const std::shared_ptr<IKernelLogic>& kernelLogic = nullptr);

ydb/core/formats/arrow/program/chain.cpp

@@ -45,8 +45,10 @@ TConclusion<TProgramChain> TProgramChain::Build(std::vector<std::shared_ptr<IRes
     THashSet<TColumnChainInfo> sourceColumns;
     std::optional<ui32> lastFilter;
     std::optional<ui32> firstAggregation;
+    std::vector<std::vector<TColumnChainInfo>> originalsToUse;
+    originalsToUse.resize(processors.size());
     for (auto&& i : processors) {
-        if (i->GetProcessorType() == EProcessorType::Aggregation) {
+        if (!firstAggregation && i->IsAggregation()) {
             firstAggregation = stepIdx;
         }
         if (!firstAggregation && i->GetProcessorType() == EProcessorType::Filter) {
@@ -62,9 +64,12 @@ TConclusion<TProgramChain> TProgramChain::Build(std::vector<std::shared_ptr<IRes
         }
         for (auto&& c : i->GetInput()) {
             auto it = contextUsage.find(c);
+            const bool isOriginalColumn = resolver.HasColumn(c);
+            if (isOriginalColumn) {
+                originalsToUse[stepIdx].emplace_back(c);
+            }
             if (it == contextUsage.end()) {
-                if (!resolver.GetColumnName(c, false)) {
-                    resolver.GetColumnName(c, true);
+                if (!isOriginalColumn) {
                     return TConclusionStatus::Fail("incorrect input column: " + ::ToString(c));
                 }
                 it = contextUsage.emplace(c, TColumnUsage::Fetch(stepIdx, i)).first;
@@ -94,7 +99,8 @@ TConclusion<TProgramChain> TProgramChain::Build(std::vector<std::shared_ptr<IRes
     }
     TProgramChain result;
     for (ui32 i = 0; i < processors.size(); ++i) {
-        result.Processors.emplace_back(std::move(columnsToFetch[i]), std::move(processors[i]), std::move(columnsToDrop[i]));
+        result.Processors.emplace_back(
+            std::move(columnsToFetch[i]), std::move(originalsToUse[i]), std::move(processors[i]), std::move(columnsToDrop[i]));
     }
     auto initStatus = result.Initialize();
     result.LastOriginalDataFilter = lastFilter;

ydb/core/formats/arrow/program/chain.h

@@ -19,6 +19,10 @@ class TProgramChain {
 public:
     TProgramChain() = default;
 
+    bool HasAggregations() const {
+        return !!FirstAggregation;
+    }
+
     bool IsGenerated(const ui32 columnId) const {
         auto it = SourcesByColumnId.find(columnId);
         AFL_VERIFY(it != SourcesByColumnId.end());

ydb/core/formats/arrow/program/filter.h

@@ -9,6 +9,10 @@ class TFilterProcessor: public IResourceProcessor {
 
     virtual TConclusionStatus DoExecute(const std::shared_ptr<TAccessorsCollection>& resources) const override;
 
+    virtual bool IsAggregation() const override {
+        return false;
+    }
+
 public:
     TFilterProcessor(std::vector<TColumnChainInfo>&& input)
         : TBase(std::move(input), {}, EProcessorType::Filter) {

ydb/core/formats/arrow/program/functions.h

@@ -26,6 +26,8 @@ class IStepFunction {
     bool NeedConcatenation = false;
 
 public:
+    virtual bool IsAggregation() const = 0;
+
     arrow::compute::ExecContext* GetContext() const {
         return GetCustomExecContext();
     }
@@ -70,6 +72,10 @@ class TSimpleFunction: public TInternalFunction {
         return { GetFunctionName(OperationId) };
     }
 
+    virtual bool IsAggregation() const override {
+        return false;
+    }
+
 public:
     static const char* GetFunctionName(const EOperation op) {
         switch (op) {
@@ -321,6 +327,10 @@ class TKernelFunction: public IStepFunction {
     const std::shared_ptr<arrow::compute::ScalarFunction> Function;
     std::shared_ptr<arrow::compute::FunctionOptions> FunctionOptions;
 
+    virtual bool IsAggregation() const override {
+        return false;
+    }
+
 public:
     TKernelFunction(const std::shared_ptr<arrow::compute::ScalarFunction> kernelsFunction,
         const std::shared_ptr<arrow::compute::FunctionOptions>& functionOptions = nullptr, const bool needConcatenation = false)

ydb/core/formats/arrow/program/kernel_logic.h

@@ -16,6 +16,8 @@ class IKernelLogic {
 
     using TFactory = NObjectFactory::TObjectFactory<IKernelLogic, TString>;
 
+    virtual TString GetClassName() const = 0;
+
     TConclusion<bool> Execute(const std::vector<TColumnChainInfo>& input, const std::vector<TColumnChainInfo>& output,
         const std::shared_ptr<TAccessorsCollection>& resources) const {
         if (!resources) {
@@ -31,6 +33,10 @@ class TGetJsonPath: public IKernelLogic {
         return "JsonValue";
     }
 private:
+    virtual TString GetClassName() const override {
+        return GetClassNameStatic();
+    }
+
     static const inline TFactory::TRegistrator<TGetJsonPath> Registrator = TFactory::TRegistrator<TGetJsonPath>(GetClassNameStatic());
 
     virtual TConclusion<bool> DoExecute(const std::vector<TColumnChainInfo>& input, const std::vector<TColumnChainInfo>& output,

ydb/core/formats/arrow/program/projection.h

@@ -9,6 +9,10 @@ class TProjectionProcessor: public IResourceProcessor {
 
     virtual TConclusionStatus DoExecute(const std::shared_ptr<TAccessorsCollection>& resources) const override;
 
+    virtual bool IsAggregation() const override {
+        return false;
+    }
+
 public:
     TProjectionProcessor(std::vector<TColumnChainInfo>&& columns)
         : TBase(std::vector<TColumnChainInfo>(columns), {}, EProcessorType::Projection) {