Optimize function transformer for NOT / IN

internal/sqlbuilder_sqlbuilder.go

@@ -72,6 +72,8 @@ const (
 	ExpressionTypeLiteral
 	ExpressionTypeParameter
 	ExpressionTypeFunction
+	ExpressionTypeList
+	ExpressionTypeUnary
 	ExpressionTypeBinary
 	ExpressionTypeSubquery
 	ExpressionTypeStar
@@ -98,6 +100,34 @@ type ExistsExpression struct {
 	Subquery *SelectStatement
 }
 
+// ListExpression represents SQL list expressions
+type ListExpression struct {
+	Expressions []*SQLExpression
+}
+
+func (e *ListExpression) WriteSql(writer *SQLWriter) {
+	writer.Write("(")
+	for i, expr := range e.Expressions {
+		writer.Write(expr.String())
+		if i != len(e.Expressions)-1 {
+			writer.Write(",")
+		}
+	}
+	writer.Write(")")
+}
+
+// UnaryExpression represents SQL unary expressions
+type UnaryExpression struct {
+	Operator   string
+	Expression *SQLExpression
+}
+
+func (e *UnaryExpression) WriteSql(writer *SQLWriter) {
+	writer.Write(e.Operator)
+	writer.Write(" ")
+	e.Expression.WriteSql(writer)
+}
+
 type BinaryExpression struct {
 	Left     *SQLExpression
 	Right    *SQLExpression
@@ -125,6 +155,8 @@ func (e *BinaryExpression) String() string {
 type SQLExpression struct {
 	Type             ExpressionType
 	Value            string
+	ListExpression   *ListExpression
+	UnaryExpression  *UnaryExpression
 	BinaryExpression *BinaryExpression
 	FunctionCall     *FunctionCall
 	Subquery         *SelectStatement
@@ -145,6 +177,10 @@ func (e *SQLExpression) WriteSql(writer *SQLWriter) {
 		}
 	case ExpressionTypeLiteral:
 		writer.Write(e.Value)
+	case ExpressionTypeList:
+		e.ListExpression.WriteSql(writer)
+	case ExpressionTypeUnary:
+		e.UnaryExpression.WriteSql(writer)
 	case ExpressionTypeBinary:
 		e.BinaryExpression.WriteSql(writer)
 	case ExpressionTypeFunction:
@@ -1105,6 +1141,27 @@ func NewSubqueryFromItem(subquery *SelectStatement, alias string) *FromItem {
 	}
 }
 
+// NewListExpression creates a new list expression
+func NewListExpression(expressions []*SQLExpression) *SQLExpression {
+	return &SQLExpression{
+		Type: ExpressionTypeList,
+		ListExpression: &ListExpression{
+			Expressions: expressions,
+		},
+	}
+}
+
+// NewNotExpression creates a new NOT expression
+func NewNotExpression(expression *SQLExpression) *SQLExpression {
+	return &SQLExpression{
+		Type: ExpressionTypeUnary,
+		UnaryExpression: &UnaryExpression{
+			Operator:   "NOT",
+			Expression: expression,
+		},
+	}
+}
+
 // NewExistsExpression creates a new EXISTS expression
 func NewExistsExpression(subquery *SelectStatement) *SQLExpression {
 	return &SQLExpression{

internal/transformer_function.go

@@ -239,6 +239,12 @@ func canOptimizeFunction(function *FunctionCallData) bool {
 
 	// Check argument count requirements
 	switch function.Name {
+	case "zetasqlite_in":
+		return true
+	case "zetasqlite_not":
+		if len(function.Arguments) != 1 {
+			return false
+		}
 	case "zetasqlite_and", "zetasqlite_or":
 		if len(function.Arguments) < 2 {
 			return false
@@ -267,9 +273,11 @@ var functionToOperator = map[string]string{
 	"zetasqlite_greater":          ">",
 	"zetasqlite_less_or_equal":    "<=",
 	"zetasqlite_greater_or_equal": ">=",
+	"zetasqlite_in":               "IN",
 	// Logical operators
 	"zetasqlite_and": "AND",
 	"zetasqlite_or":  "OR",
+	"zetasqlite_not": "NOT",
 }
 
 // optimizeFunctionToSQL converts functions to direct SQL operators
@@ -291,6 +299,13 @@ func optimizeFunctionToSQL(functionName string, args []*SQLExpression) (*SQLExpr
 		}
 		return result, nil
 
+	case "zetasqlite_not":
+		if len(args) != 1 {
+			return nil, fmt.Errorf("%s expected only 1 argument, got %d", functionName, len(args))
+		}
+		return NewNotExpression(args[0]), nil
+	case "zetasqlite_in":
+		return NewBinaryExpression(args[0], operator, NewListExpression(args[1:])), nil
 	default: // comparison operators
 		if len(args) != 2 {
 			return nil, fmt.Errorf("%s expected 2 arguments, got %d", functionName, len(args))

query_test.go

@@ -16,7 +16,7 @@
 )
 
 func TestQuery(t *testing.T) {
 	os.Setenv("TZ", "UTC")
 	now := time.Now()
 	ctx := context.Background()
 	ctx = zetasqlite.WithCurrentTime(ctx, now)
@@ -6109,16 +6109,16 @@
 		{
 			name: "function call with many arguments - under limit",
 			query: fmt.Sprintf(
-				`select %s true in (true)`,
-				strings.Repeat("false in (true) or ", 999),
+				`select concat(%s "a")`,
+				strings.Repeat(`"a",`, 998),
 			),
-			expectedRows: [][]interface{}{{true}},
+			expectedRows: [][]interface{}{{strings.Repeat("a", 999)}},
 		},
 		{
 			name: "function call with many arguments - above limit",
 			query: fmt.Sprintf(
-				`select %s true in (true)`,
-				strings.Repeat("false in (true) or ", 1001),
+				`select concat(%s "a")`,
+				strings.Repeat(`"a",`, 1000),
 			),
 			expectedErr: "too many arguments on function",
 		},
@@ -6301,7 +6301,7 @@
 			}
 		})
 	}
 	os.Unsetenv("TZ")
 }
 
 func createTimestampFormatFromTime(t time.Time) string {