filter chains optimization

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

@@ -20,13 +20,13 @@ NJson::TJsonValue IResourceProcessor::DebugJson() const {
     return result;
 }
 
-TConclusionStatus IResourceProcessor::Execute(const std::shared_ptr<TAccessorsCollection>& resources) const {
+TConclusionStatus IResourceProcessor::Execute(const std::shared_ptr<TAccessorsCollection>& resources, const TProcessorContext& context) const {
     for (auto&& i : Output) {
         if (resources->HasColumn(i.GetColumnId())) {
             return TConclusionStatus::Fail("column " + ::ToString(i.GetColumnId()) + " has already");
         }
     }
-    return DoExecute(resources);
+    return DoExecute(resources, context);
 }
 
 std::optional<TFetchingInfo> IResourceProcessor::BuildFetchTask(

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

@@ -164,7 +164,8 @@ enum class EProcessorType {
     Calculation,
     Projection,
     Filter,
-    Aggregation
+    Aggregation,
+    Original
 };
 
 class TFetchingInfo {
@@ -187,20 +188,53 @@ class TFetchingInfo {
     }
 };
 
+class TProcessorContext {
+protected:
+    std::vector<TColumnChainInfo> ColumnsToFetch;
+    std::vector<TColumnChainInfo> OriginalColumnsToUse;
+    std::vector<TColumnChainInfo> ColumnsToDrop;
+
+public:
+    const std::vector<TColumnChainInfo>& GetColumnsToFetch() const {
+        return ColumnsToFetch;
+    }
+    const std::vector<TColumnChainInfo>& GetOriginalColumnsToUse() const {
+        return OriginalColumnsToUse;
+    }
+    const std::vector<TColumnChainInfo>& GetColumnsToDrop() const {
+        return ColumnsToDrop;
+    }
+
+    TProcessorContext(
+        std::vector<TColumnChainInfo>&& toFetch, std::vector<TColumnChainInfo>&& originalToUse, std::vector<TColumnChainInfo>&& toDrop)
+        : ColumnsToFetch(std::move(toFetch))
+        , OriginalColumnsToUse(std::move(originalToUse))
+        , ColumnsToDrop(std::move(toDrop)) {
+    }
+};
+
 class IResourceProcessor {
 private:
     YDB_READONLY_DEF(std::vector<TColumnChainInfo>, Input);
     YDB_READONLY_DEF(std::vector<TColumnChainInfo>, Output);
     YDB_READONLY(EProcessorType, ProcessorType, EProcessorType::Unknown);
 
-    virtual TConclusionStatus DoExecute(const std::shared_ptr<TAccessorsCollection>& resources) const = 0;
+    virtual TConclusionStatus DoExecute(const std::shared_ptr<TAccessorsCollection>& resources, const TProcessorContext& context) const = 0;
 
     virtual NJson::TJsonValue DoDebugJson() const {
         return NJson::JSON_MAP;
     }
+    virtual ui64 DoGetWeight() const {
+        return 0;
+    }
 
 public:
-    virtual std::optional<TFetchingInfo> BuildFetchTask(const ui32 columnId, const NAccessor::IChunkedArray::EType arrType, const std::shared_ptr<TAccessorsCollection>& resources) const;
+    ui64 GetWeight() const {
+        return DoGetWeight();
+    }
+
+    virtual std::optional<TFetchingInfo> BuildFetchTask(
+        const ui32 columnId, const NAccessor::IChunkedArray::EType arrType, const std::shared_ptr<TAccessorsCollection>& resources) const;
 
     virtual bool IsAggregation() const = 0;
 
@@ -224,25 +258,21 @@ class IResourceProcessor {
         , ProcessorType(type) {
     }
 
-    [[nodiscard]] TConclusionStatus Execute(const std::shared_ptr<TAccessorsCollection>& resources) const;
+    [[nodiscard]] TConclusionStatus Execute(const std::shared_ptr<TAccessorsCollection>& resources, const TProcessorContext& context) const;
 };
 
-class TResourceProcessorStep {
+class TResourceProcessorStep: public TProcessorContext {
 private:
-    YDB_READONLY_DEF(std::vector<TColumnChainInfo>, ColumnsToFetch);
-    YDB_READONLY_DEF(std::vector<TColumnChainInfo>, OriginalColumnsToUse);
+    using TBase = TProcessorContext;
     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::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)) {
+        : TBase(std::move(toFetch), std::move(originalToUse), std::move(toDrop))
+        , Processor(std::move(processor)) {
         AFL_VERIFY(Processor);
     }
 

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

@@ -66,7 +66,8 @@ CH::AggFunctionId TWithKeysAggregationOption::GetHouseFunction(const EAggregate
     return CH::AggFunctionId::AGG_UNSPECIFIED;
 }
 
-TConclusionStatus TWithKeysAggregationProcessor::DoExecute(const std::shared_ptr<TAccessorsCollection>& resources) const {
+TConclusionStatus TWithKeysAggregationProcessor::DoExecute(
+    const std::shared_ptr<TAccessorsCollection>& resources, const TProcessorContext& /*context*/) const {
     CH::GroupByOptions funcOpts;
     funcOpts.assigns.reserve(AggregationKeys.size() + Aggregations.size());
     funcOpts.has_nullable_key = false;

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

@@ -20,7 +20,7 @@ class TAggregateFunction: public TInternalFunction {
     }
     virtual TConclusion<arrow::Datum> Call(
         const TExecFunctionContext& context, const std::shared_ptr<TAccessorsCollection>& resources) const override;
-    
+
     TConclusion<arrow::Datum> PrepareResult(arrow::Datum&& datum) const override {
         if (!datum.is_scalar()) {
             return TConclusionStatus::Fail("Aggregate result is not a scalar.");
@@ -103,13 +103,14 @@ class TAggregateFunction: public TInternalFunction {
         return "";
     }
 
-    virtual TConclusionStatus CheckIO(const std::vector<TColumnChainInfo>& /*input*/, const std::vector<TColumnChainInfo>& output) const override {
+    virtual TConclusionStatus CheckIO(
+        const std::vector<TColumnChainInfo>& /*input*/, const std::vector<TColumnChainInfo>& output) const override {
         if (output.size() != 1) {
             return TConclusionStatus::Fail("output size != 1 (" + ::ToString(output.size()) + ")");
         }
-//        if (input.size() != 1) {
-//            return TConclusionStatus::Fail("input size != 1 (" + ::ToString(input.size()) + ")");
-//        }
+        //        if (input.size() != 1) {
+        //            return TConclusionStatus::Fail("input size != 1 (" + ::ToString(input.size()) + ")");
+        //        }
         return TConclusionStatus::Success();
     }
 };
@@ -149,7 +150,7 @@ class TWithKeysAggregationProcessor: public IResourceProcessor {
     std::vector<TColumnChainInfo> AggregationKeys;
     std::vector<TWithKeysAggregationOption> Aggregations;
 
-    virtual TConclusionStatus DoExecute(const std::shared_ptr<TAccessorsCollection>& resources) const override;
+    virtual TConclusionStatus DoExecute(const std::shared_ptr<TAccessorsCollection>& resources, const TProcessorContext& context) const override;
 
     TWithKeysAggregationProcessor(std::vector<TColumnChainInfo>&& input, std::vector<TColumnChainInfo>&& output,
         std::vector<TColumnChainInfo>&& aggregationKeys, std::vector<TWithKeysAggregationOption>&& aggregations)

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

@@ -10,7 +10,7 @@
 
 namespace NKikimr::NArrow::NSSA {
 
-TConclusionStatus TConstProcessor::DoExecute(const std::shared_ptr<TAccessorsCollection>& resources) const {
+TConclusionStatus TConstProcessor::DoExecute(const std::shared_ptr<TAccessorsCollection>& resources, const TProcessorContext& /*context*/) const {
     AFL_VERIFY(GetInput().empty());
     resources->AddConstantVerified(GetOutputColumnIdOnce(), ScalarConstant);
     return TConclusionStatus::Success();

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

@@ -8,7 +8,7 @@ class TConstProcessor: public IResourceProcessor {
     using TBase = IResourceProcessor;
     YDB_READONLY_DEF(std::shared_ptr<arrow::Scalar>, ScalarConstant);
 
-    virtual TConclusionStatus DoExecute(const std::shared_ptr<TAccessorsCollection>& resources) const override;
+    virtual TConclusionStatus DoExecute(const std::shared_ptr<TAccessorsCollection>& resources, const TProcessorContext& context) const override;
 
     virtual bool IsAggregation() const override {
         return false;

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

@@ -4,7 +4,8 @@
 
 namespace NKikimr::NArrow::NSSA {
 
-TConclusionStatus TCalculationProcessor::DoExecute(const std::shared_ptr<TAccessorsCollection>& resources) const {
+TConclusionStatus TCalculationProcessor::DoExecute(
+    const std::shared_ptr<TAccessorsCollection>& resources, const TProcessorContext& /*context*/) const {
     if (KernelLogic) {
         auto resultKernel = KernelLogic->Execute(GetInput(), GetOutput(), resources);
         if (resultKernel.IsFail()) {

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

@@ -3,6 +3,8 @@
 #include "functions.h"
 #include "kernel_logic.h"
 
+#include <yql/essentials/core/arrow_kernels/request/request.h>
+
 namespace NKikimr::NArrow::NSSA {
 
 class TCalculationProcessor: public IResourceProcessor {
@@ -14,7 +16,7 @@ class TCalculationProcessor: public IResourceProcessor {
 
     std::shared_ptr<IStepFunction> Function;
 
-    virtual TConclusionStatus DoExecute(const std::shared_ptr<TAccessorsCollection>& resources) const override;
+    virtual TConclusionStatus DoExecute(const std::shared_ptr<TAccessorsCollection>& resources, const TProcessorContext& context) const override;
 
     TCalculationProcessor(std::vector<TColumnChainInfo>&& input, std::vector<TColumnChainInfo>&& output,
         const std::shared_ptr<IStepFunction>& function, const std::shared_ptr<IKernelLogic>& kernelLogic)
@@ -27,6 +29,23 @@ class TCalculationProcessor: public IResourceProcessor {
         return Function->IsAggregation();
     }
 
+    virtual ui64 DoGetWeight() const override {
+        if (KernelLogic) {
+            return 0;
+        }
+        if (!YqlOperationId) {
+            return 10;
+        } else if ((NYql::TKernelRequestBuilder::EBinaryOp)*YqlOperationId == NYql::TKernelRequestBuilder::EBinaryOp::StartsWith ||
+                   (NYql::TKernelRequestBuilder::EBinaryOp)*YqlOperationId == NYql::TKernelRequestBuilder::EBinaryOp::EndsWith) {
+            return 7;
+        } else if ((NYql::TKernelRequestBuilder::EBinaryOp)*YqlOperationId == NYql::TKernelRequestBuilder::EBinaryOp::StringContains) {
+            return 10;
+        } else if ((NYql::TKernelRequestBuilder::EBinaryOp)*YqlOperationId == NYql::TKernelRequestBuilder::EBinaryOp::Equals) {
+            return 5;
+        }
+        return 0;
+    }
+
 public:
     virtual std::optional<TFetchingInfo> BuildFetchTask(const ui32 columnId, const NAccessor::IChunkedArray::EType arrType,
         const std::shared_ptr<TAccessorsCollection>& resources) const override {

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

@@ -1,5 +1,6 @@
 #include "chain.h"
 #include "collection.h"
+#include "graph.h"
 
 namespace NKikimr::NArrow::NSSA {
 
@@ -39,7 +40,17 @@ class TColumnUsage {
 };
 }   // namespace
 
-TConclusion<TProgramChain> TProgramChain::Build(std::vector<std::shared_ptr<IResourceProcessor>>&& processors, const IColumnResolver& resolver) {
+TConclusion<TProgramChain> TProgramChain::Build(std::vector<std::shared_ptr<IResourceProcessor>>&& processorsExt, const IColumnResolver& resolver) {
+    NOptimization::TGraph graph(std::move(processorsExt), resolver);
+    auto conclusion = graph.Collapse();
+    if (conclusion.IsFail()) {
+        return conclusion;
+    }
+    auto processorsConclusion = graph.BuildChain();
+    if (processorsConclusion.IsFail()) {
+        return processorsConclusion;
+    }
+    auto processors = processorsConclusion.DetachResult();
     THashMap<ui32, TColumnUsage> contextUsage;
     ui32 stepIdx = 0;
     THashSet<ui32> sourceColumns;
@@ -149,7 +160,7 @@ TConclusionStatus TProgramChain::Initialize() {
 
 TConclusionStatus TProgramChain::Apply(const std::shared_ptr<TAccessorsCollection>& resources) const {
     for (auto&& i : Processors) {
-        auto status = i->Execute(resources);
+        auto status = i->Execute(resources, i);
         if (status.IsFail()) {
             return status;
         }

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

@@ -60,21 +60,21 @@ class TFilterVisitor: public arrow::ArrayVisitor {
     arrow::Status VisitImpl(const TArray& array) {
         AFL_VERIFY(Started);
         for (ui32 i = 0; i < array.length(); ++i) {
-            const bool columnValue = (bool)array.Value(i);
             const ui32 currentIdx = CursorIdx++;
-            FiltersMerged[currentIdx] = FiltersMerged[currentIdx] && columnValue;
+            FiltersMerged[currentIdx] = FiltersMerged[currentIdx] && !array.IsNull(i) && (bool)array.Value(i);
         }
         AFL_VERIFY(CursorIdx <= FiltersMerged.size());
         return arrow::Status::OK();
     }
 };
 
-TConclusionStatus TFilterProcessor::DoExecute(const std::shared_ptr<TAccessorsCollection>& resources) const {
+TConclusionStatus TFilterProcessor::DoExecute(const std::shared_ptr<TAccessorsCollection>& resources, const TProcessorContext& context) const {
     std::vector<std::shared_ptr<IChunkedArray>> inputColumns;
-    if (ReuseColumns) {
-        inputColumns = resources->GetAccessors(TColumnChainInfo::ExtractColumnIds(GetInput()));
-    } else {
+    AFL_VERIFY(context.GetColumnsToDrop().size() <= 1)("size", context.GetColumnsToDrop().size());
+    if (context.GetColumnsToDrop().size() && GetInputColumnIdOnce() == context.GetColumnsToDrop().front()) {
         inputColumns = resources->ExtractAccessors(TColumnChainInfo::ExtractColumnIds(GetInput()));
+    } else {
+        inputColumns = resources->GetAccessors(TColumnChainInfo::ExtractColumnIds(GetInput()));
     }
     TFilterVisitor filterVisitor(inputColumns.front()->GetRecordsCount());
     for (auto& arr : inputColumns) {

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

@@ -6,25 +6,20 @@ namespace NKikimr::NArrow::NSSA {
 class TFilterProcessor: public IResourceProcessor {
 private:
     using TBase = IResourceProcessor;
-    const bool ReuseColumns;
-    virtual TConclusionStatus DoExecute(const std::shared_ptr<TAccessorsCollection>& resources) const override;
+    virtual TConclusionStatus DoExecute(const std::shared_ptr<TAccessorsCollection>& resources, const TProcessorContext& context) const override;
 
     virtual bool IsAggregation() const override {
         return false;
     }
 
 public:
-    TFilterProcessor(std::vector<TColumnChainInfo>&& input, const bool reuseColumns = false)
-        : TBase(std::move(input), {}, EProcessorType::Filter)
-        , ReuseColumns(reuseColumns)
-    {
-        AFL_VERIFY(GetInput().size());
+    TFilterProcessor(std::vector<TColumnChainInfo>&& input)
+        : TBase(std::move(input), {}, EProcessorType::Filter) {
+        AFL_VERIFY(GetInput().size() == 1)("size", GetInput().size());
     }
 
-    TFilterProcessor(const TColumnChainInfo& input, const bool reuseColumns = false)
-        : TBase({ input }, {}, EProcessorType::Filter)
-        , ReuseColumns(reuseColumns)
-    {
+    TFilterProcessor(const TColumnChainInfo& input)
+        : TBase({ input }, {}, EProcessorType::Filter) {
     }
 };