Skip to content

Commit 7105766

Browse files
yhmoyhmo
committed
Support FunctionChain
Signed-off-by: yhmo <yihua.mo@zilliz.com>
1 parent 382eb27 commit 7105766

10 files changed

Lines changed: 978 additions & 2 deletions

File tree

‎sdk-core/src/main/java/io/milvus/v2/service/vector/request/FunctionChain.java‎

Lines changed: 130 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,130 @@
1+
/*
2+
* Licensed to the Apache Software Foundation (ASF) under one
3+
* or more contributor license agreements. See the NOTICE file
4+
* distributed with this work for additional information
5+
* regarding copyright ownership. The ASF licenses this file
6+
* to you under the Apache License, Version 2.0 (the
7+
* "License"); you may not use this file except in compliance
8+
* with the License. You may obtain a copy of the License at
9+
*
10+
* http://www.apache.org/licenses/LICENSE-2.0
11+
*
12+
* Unless required by applicable law or agreed to in writing,
13+
* software distributed under the License is distributed on an
14+
* "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
15+
* KIND, either express or implied. See the License for the
16+
* specific language governing permissions and limitations
17+
* under the License.
18+
*/
19+
20+
package io.milvus.v2.service.vector.request;
21+
22+
import io.milvus.v2.exception.ErrorCode;
23+
import io.milvus.v2.exception.MilvusClientException;
24+
25+
import java.util.ArrayList;
26+
import java.util.List;
27+
28+
/**
29+
* An ordered rerank/refine plan applied to search results. Mirrors PyMilvus's
30+
* {@code FunctionChain}: a fluent builder that composes {@code map}, {@code sort}, and
31+
* {@code limit} operations and serializes to the gRPC {@code FunctionChain} message.
32+
*
33+
* <pre>{@code
34+
* FunctionChain chain = FunctionChain.builder()
35+
* .stage(FunctionChainStage.L2_RERANK)
36+
* .name("fresh_popular_rerank")
37+
* .map("$score", FunctionChainExpr.builder()
38+
* .name("num_combine")
39+
* .arg(FunctionChainArg.col("$score"))
40+
* .arg(FunctionChainArg.col("freshness"))
41+
* .param("mode", "weighted")
42+
* .param("weights", Arrays.asList(0.7, 0.2, 0.1))
43+
* .build())
44+
* .sort("$score", true, "$id")
45+
* .limit(10)
46+
* .build();
47+
* }</pre>
48+
*/
49+
public class FunctionChain {
50+
private final FunctionChainStage stage;
51+
private final String name;
52+
private final List<FunctionChainOp> ops;
53+
54+
private FunctionChain(FunctionChainBuilder builder) {
55+
this.stage = builder.stage;
56+
this.name = builder.name;
57+
this.ops = new ArrayList<>(builder.ops);
58+
}
59+
60+
public static FunctionChainBuilder builder() {
61+
return new FunctionChainBuilder();
62+
}
63+
64+
public FunctionChainStage getStage() {
65+
return stage;
66+
}
67+
68+
public String getName() {
69+
return name;
70+
}
71+
72+
public List<FunctionChainOp> getOps() {
73+
return ops;
74+
}
75+
76+
public io.milvus.grpc.FunctionChain toGrpc() {
77+
io.milvus.grpc.FunctionChain.Builder builder = io.milvus.grpc.FunctionChain.newBuilder()
78+
.setName(name)
79+
.setStage(stage.toGrpc());
80+
ops.forEach(op -> builder.addOps(op.toGrpc()));
81+
return builder.build();
82+
}
83+
84+
public static class FunctionChainBuilder {
85+
private FunctionChainStage stage = FunctionChainStage.UNSPECIFIED;
86+
private String name = "";
87+
private final List<FunctionChainOp> ops = new ArrayList<>();
88+
89+
private FunctionChainBuilder() {
90+
}
91+
92+
public FunctionChainBuilder stage(FunctionChainStage stage) {
93+
if (stage == null) {
94+
throw new MilvusClientException(ErrorCode.INVALID_PARAMS,
95+
"Function chain stage must not be null");
96+
}
97+
this.stage = stage;
98+
return this;
99+
}
100+
101+
public FunctionChainBuilder name(String name) {
102+
this.name = name == null ? "" : name;
103+
return this;
104+
}
105+
106+
public FunctionChainBuilder map(String output, FunctionChainExpr expr) {
107+
this.ops.add(FunctionChainOp.map(output, expr));
108+
return this;
109+
}
110+
111+
public FunctionChainBuilder sort(String by, boolean desc, String tieBreakCol) {
112+
this.ops.add(FunctionChainOp.sort(by, desc, tieBreakCol));
113+
return this;
114+
}
115+
116+
public FunctionChainBuilder limit(int limit) {
117+
this.ops.add(FunctionChainOp.limit(limit, 0));
118+
return this;
119+
}
120+
121+
public FunctionChainBuilder limit(int limit, int offset) {
122+
this.ops.add(FunctionChainOp.limit(limit, offset));
123+
return this;
124+
}
125+
126+
public FunctionChain build() {
127+
return new FunctionChain(this);
128+
}
129+
}
130+
}

‎sdk-core/src/main/java/io/milvus/v2/service/vector/request/FunctionChainArg.java‎

Lines changed: 71 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,71 @@
1+
/*
2+
* Licensed to the Apache Software Foundation (ASF) under one
3+
* or more contributor license agreements. See the NOTICE file
4+
* distributed with this work for additional information
5+
* regarding copyright ownership. The ASF licenses this file
6+
* to you under the Apache License, Version 2.0 (the
7+
* "License"); you may not use this file except in compliance
8+
* with the License. You may obtain a copy of the License at
9+
*
10+
* http://www.apache.org/licenses/LICENSE-2.0
11+
*
12+
* Unless required by applicable law or agreed to in writing,
13+
* software distributed under the License is distributed on an
14+
* "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
15+
* KIND, either express or implied. See the License for the
16+
* specific language governing permissions and limitations
17+
* under the License.
18+
*/
19+
20+
package io.milvus.v2.service.vector.request;
21+
22+
import io.milvus.v2.exception.ErrorCode;
23+
import io.milvus.v2.exception.MilvusClientException;
24+
25+
/**
26+
* An argument of a function-chain expression: either a collection-field reference or a literal
27+
* value. Mirrors PyMilvus's {@code FunctionChainArg} union of {@code ColumnRef} and literal.
28+
*/
29+
public class FunctionChainArg {
30+
private final String columnName;
31+
private final FunctionParamValue literal;
32+
33+
private FunctionChainArg(String columnName, FunctionParamValue literal) {
34+
this.columnName = columnName;
35+
this.literal = literal;
36+
}
37+
38+
public static FunctionChainArg col(String name) {
39+
if (name == null || name.isEmpty()) {
40+
throw new MilvusClientException(ErrorCode.INVALID_PARAMS, "Column name must not be empty");
41+
}
42+
return new FunctionChainArg(name, null);
43+
}
44+
45+
public static FunctionChainArg literal(Object value) {
46+
return new FunctionChainArg(null, FunctionParamValue.from(value));
47+
}
48+
49+
public boolean isColumn() {
50+
return columnName != null;
51+
}
52+
53+
public String getColumnName() {
54+
return columnName;
55+
}
56+
57+
public FunctionParamValue getLiteral() {
58+
return literal;
59+
}
60+
61+
public io.milvus.grpc.FunctionChainExprArg toGrpc() {
62+
if (columnName != null) {
63+
return io.milvus.grpc.FunctionChainExprArg.newBuilder()
64+
.setColumn(io.milvus.grpc.FunctionChainColumnArg.newBuilder().setName(columnName).build())
65+
.build();
66+
}
67+
return io.milvus.grpc.FunctionChainExprArg.newBuilder()
68+
.setLiteral(literal.toGrpc())
69+
.build();
70+
}
71+
}

‎sdk-core/src/main/java/io/milvus/v2/service/vector/request/FunctionChainExpr.java‎

Lines changed: 110 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,110 @@
1+
/*
2+
* Licensed to the Apache Software Foundation (ASF) under one
3+
* or more contributor license agreements. See the NOTICE file
4+
* distributed with this work for additional information
5+
* regarding copyright ownership. The ASF licenses this file
6+
* to you under the Apache License, Version 2.0 (the
7+
* "License"); you may not use this file except in compliance
8+
* with the License. You may obtain a copy of the License at
9+
*
10+
* http://www.apache.org/licenses/LICENSE-2.0
11+
*
12+
* Unless required by applicable law or agreed to in writing,
13+
* software distributed under the License is distributed on an
14+
* "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
15+
* KIND, either express or implied. See the License for the
16+
* specific language governing permissions and limitations
17+
* under the License.
18+
*/
19+
20+
package io.milvus.v2.service.vector.request;
21+
22+
import io.milvus.v2.exception.ErrorCode;
23+
import io.milvus.v2.exception.MilvusClientException;
24+
25+
import java.util.ArrayList;
26+
import java.util.LinkedHashMap;
27+
import java.util.List;
28+
import java.util.Map;
29+
30+
/**
31+
* A named function invocation used by a function-chain {@code map} operation.
32+
*
33+
* <p>Arguments are added with {@code arg(FunctionChainArg.col(...))} or
34+
* {@code arg(FunctionChainArg.literal(...))}, and keyword parameters with
35+
* {@code param(String, Object)}. This mirrors PyMilvus's {@code FunctionChainExpr}.
36+
*/
37+
public class FunctionChainExpr {
38+
private final String name;
39+
private final List<FunctionChainArg> args;
40+
private final Map<String, FunctionParamValue> params;
41+
42+
private FunctionChainExpr(FunctionChainExprBuilder builder) {
43+
this.name = builder.name;
44+
this.args = new ArrayList<>(builder.args);
45+
this.params = new LinkedHashMap<>(builder.params);
46+
}
47+
48+
public static FunctionChainExprBuilder builder() {
49+
return new FunctionChainExprBuilder();
50+
}
51+
52+
public String getName() {
53+
return name;
54+
}
55+
56+
public List<FunctionChainArg> getArgs() {
57+
return args;
58+
}
59+
60+
public Map<String, FunctionParamValue> getParams() {
61+
return params;
62+
}
63+
64+
public io.milvus.grpc.FunctionChainExpr toGrpc() {
65+
io.milvus.grpc.FunctionChainExpr.Builder builder = io.milvus.grpc.FunctionChainExpr.newBuilder().setName(name);
66+
args.forEach(arg -> builder.addArgs(arg.toGrpc()));
67+
params.forEach((k, v) -> builder.putParams(k, v.toGrpc()));
68+
return builder.build();
69+
}
70+
71+
public static class FunctionChainExprBuilder {
72+
private String name;
73+
private final List<FunctionChainArg> args = new ArrayList<>();
74+
private final Map<String, FunctionParamValue> params = new LinkedHashMap<>();
75+
76+
private FunctionChainExprBuilder() {
77+
}
78+
79+
public FunctionChainExprBuilder name(String name) {
80+
this.name = name;
81+
return this;
82+
}
83+
84+
public FunctionChainExprBuilder arg(FunctionChainArg arg) {
85+
if (arg == null) {
86+
throw new MilvusClientException(ErrorCode.INVALID_PARAMS,
87+
"Function chain expression arg must not be null");
88+
}
89+
this.args.add(arg);
90+
return this;
91+
}
92+
93+
public FunctionChainExprBuilder param(String key, Object value) {
94+
if (key == null || key.isEmpty()) {
95+
throw new MilvusClientException(ErrorCode.INVALID_PARAMS,
96+
"Function chain expression parameter names must be non-empty strings");
97+
}
98+
this.params.put(key, FunctionParamValue.from(value));
99+
return this;
100+
}
101+
102+
public FunctionChainExpr build() {
103+
if (name == null || name.isEmpty()) {
104+
throw new MilvusClientException(ErrorCode.INVALID_PARAMS,
105+
"Function chain expression name must be a non-empty string");
106+
}
107+
return new FunctionChainExpr(this);
108+
}
109+
}
110+
}

0 commit comments

Comments
 (0)