implemented list_where, and debug a typecheck error for list_select

Author: Jasonlmx

src/function/function_collection.cpp

@@ -131,6 +131,7 @@ FunctionCollection* FunctionCollection::getFunctions() {
         SCALAR_FUNCTION(ListCosineSimilarityFunction), SCALAR_FUNCTION(ListCosineDistanceFunction),
         SCALAR_FUNCTION(ListDistanceFunction), SCALAR_FUNCTION(ListHasAnyFunction),
         SCALAR_FUNCTION(ListIntersectFunction), SCALAR_FUNCTION(ListSelectFunction),
+        SCALAR_FUNCTION(ListWhereFunction),
 
         // Cast functions
         SCALAR_FUNCTION(CastToDateFunction), SCALAR_FUNCTION_ALIAS(DateFunction),

src/function/list/CMakeLists.txt

@@ -30,7 +30,8 @@ add_library(kuzu_list_function
         list_has_all.cpp
         list_has_any.cpp
         list_intersect.cpp
-        list_select_function.cpp)
+        list_select_function.cpp
+        list_where_function.cpp)
 
 set(ALL_OBJECT_FILES
         ${ALL_OBJECT_FILES} $<TARGET_OBJECTS:kuzu_list_function>

src/function/list/list_select_function.cpp

@@ -39,9 +39,10 @@ static std::unique_ptr<FunctionBindData> bindFunc(const ScalarBindFuncInput& inp
     if (types[1].getPhysicalType()!=PhysicalTypeID::LIST) {
         throw BinderException(ExceptionMessage::listFunctionIncompatibleChildrenType(
             ListIntersectFunction::name, types[0].toString(), types[1].toString()));
+    } else {
         auto thisExtraTypeInfo=types[1].getExtraTypeInfo();
         auto thisListTypeInfo=ku_dynamic_cast<const ListTypeInfo*>(thisExtraTypeInfo);
-        if (thisListTypeInfo->getChildType().getLogicalTypeID()!=LogicalTypeID::INT64) {
+        if (thisListTypeInfo->getChildType().getPhysicalType()!=PhysicalTypeID::INT64) {
             throw BinderException("LIST_SELECT expecting argument type: LIST of ANY, LIST of INT");
         }
     }

src/function/list/list_where_function.cpp

@@ -0,0 +1,77 @@
+#include "common/exception/binder.h"
+#include "common/exception/message.h"
+#include "common/type_utils.h"
+#include "common/types/types.h"
+#include "function/list/functions/list_function_utils.h"
+#include "function/list/functions/list_position_function.h"
+#include "function/list/functions/list_unique_function.h"
+#include "function/list/vector_list_functions.h"
+#include "function/scalar_function.h"
+
+using namespace kuzu::common;
+
+namespace kuzu {
+namespace function {
+
+struct ListWhere {
+    static void operation(common::list_entry_t& left, common::list_entry_t& right,
+        common::list_entry_t& result, common::ValueVector& leftVector,
+        common::ValueVector& rightVector, common::ValueVector& resultVector) {
+        if (right.size!=left.size) {
+            throw BinderException(stringFormat("LIST_WHERE expecting lists of same size, receiving size {} and size {}", left.size, left.size));
+        }
+        auto leftDataVector = common::ListVector::getDataVector(&leftVector);
+        auto leftPos = left.offset;
+        auto rightDataVector = common::ListVector::getDataVector(&rightVector);
+        auto rightPos = right.offset;
+        list_size_t resultSize=0;
+        std::vector<bool> maskListBools;
+        for (auto i=0u; i < right.size; i++) {
+            auto maskBool=rightDataVector->getValue<bool>(rightPos+i);
+            maskListBools.push_back(maskBool);
+            if (maskBool) {
+                resultSize++;
+            }
+        }
+        result = common::ListVector::addList(&resultVector, resultSize);
+        auto resultDataVector = common::ListVector::getDataVector(&resultVector);
+        auto resultPos = result.offset;
+        for (auto i=0u; i < right.size; i++) {
+            auto maskBool=maskListBools.at(i);
+            if (maskBool) {
+                resultDataVector->copyFromVectorData(resultPos++, leftDataVector, leftPos+i);
+            }
+        }
+    }
+};
+static std::unique_ptr<FunctionBindData> bindFunc(const ScalarBindFuncInput& input) {
+    std::vector<LogicalType> types;
+    types.push_back(input.arguments[0]->getDataType().copy());
+    types.push_back(input.arguments[1]->getDataType().copy());
+    if (types[1].getPhysicalType()!=PhysicalTypeID::LIST) {
+        throw BinderException(ExceptionMessage::listFunctionIncompatibleChildrenType(
+            ListIntersectFunction::name, types[0].toString(), types[1].toString()));
+    } else {
+        auto thisExtraTypeInfo=types[1].getExtraTypeInfo();
+        auto thisListTypeInfo=ku_dynamic_cast<const ListTypeInfo*>(thisExtraTypeInfo);
+        if (thisListTypeInfo->getChildType().getPhysicalType()!=PhysicalTypeID::BOOL) {
+            throw BinderException("LIST_SELECT expecting argument type: LIST of ANY, LIST of BOOL");
+        }
+    }
+    return std::make_unique<FunctionBindData>(std::move(types), types[0].copy());
+}
+
+function_set ListWhereFunction::getFunctionSet() {
+    function_set result;
+    auto execFunc = ScalarFunction::BinaryExecListStructFunction<list_entry_t, list_entry_t,
+        list_entry_t, ListWhere>;
+    auto function = std::make_unique<ScalarFunction>(name,
+        std::vector<LogicalTypeID>{LogicalTypeID::LIST, LogicalTypeID::LIST}, LogicalTypeID::LIST,
+        execFunc);
+    function->bindFunc = bindFunc;
+    result.push_back(std::move(function));
+    return result;
+}
+
+} // namespace function
+} // namespace kuzu

src/include/function/list/vector_list_functions.h

@@ -248,5 +248,11 @@ struct ListSelectFunction {
     static function_set getFunctionSet();
 };
 
+struct ListWhereFunction {
+    static constexpr const char* name = "LIST_WHERE";
+
+    static function_set getFunctionSet();
+};
+
 } // namespace function
 } // namespace kuzu