Add stored procedure calls with unnamed args

examples/src/kotlin/org/partiql/examples/CustomProceduresExample.kt

@@ -0,0 +1,130 @@
+package org.partiql.examples
+
+import com.amazon.ion.IonDecimal
+import com.amazon.ion.IonStruct
+import com.amazon.ion.system.IonSystemBuilder
+import org.partiql.examples.util.Example
+import org.partiql.lang.CompilerPipeline
+import org.partiql.lang.errors.ErrorCode
+import org.partiql.lang.errors.Property
+import org.partiql.lang.errors.PropertyValueMap
+import org.partiql.lang.eval.BindingCase
+import org.partiql.lang.eval.BindingName
+import org.partiql.lang.eval.Bindings
+import org.partiql.lang.eval.EvaluationException
+import org.partiql.lang.eval.EvaluationSession
+import org.partiql.lang.eval.ExprValue
+import org.partiql.lang.eval.ExprValueFactory
+import org.partiql.lang.eval.ExprValueType
+import org.partiql.lang.eval.builtins.storedprocedure.StoredProcedure
+import org.partiql.lang.eval.builtins.storedprocedure.StoredProcedureSignature
+import org.partiql.lang.eval.stringValue
+import java.io.PrintStream
+import java.math.BigDecimal
+import java.math.RoundingMode
+
+private val ion = IonSystemBuilder.standard().build()
+
+/**
+ * A simple custom stored procedure that calculates the moon weight for each crewmate of the given crew, storing the
+ * moon weight in the [EvaluationSession] global bindings. This procedure also returns the number of crewmates we
+ * calculated the moon weight for, returning -1 if no crew is found.
+ *
+ * This example demonstrates how to create a custom stored procedure, check argument types, and modify the
+ * [EvaluationSession].
+ */
+class CalculateCrewMoonWeight(private val valueFactory: ExprValueFactory): StoredProcedure {
+    private val MOON_GRAVITATIONAL_CONSTANT = BigDecimal(1.622 / 9.81)
+
+    // [StoredProcedureSignature] takes two arguments:
+    //   1. the name of the stored procedure
+    //   2. the arity of this stored procedure. Checks to arity are taken care of by the evaluator. However, we must
+    //      still check that the passed arguments are of the right type in our implementation of the procedure.
+    override val signature = StoredProcedureSignature(name = "calculate_crew_moon_weight", arity = 1)
+
+    // `call` is where you define the logic of the stored procedure given an [EvaluationSession] and a list of
+    // arguments
+    override fun call(session: EvaluationSession, args: List<ExprValue>): ExprValue {
+        // We first check that the first argument is a string
+        val crewName = args.first()
+        // In the future the evaluator will also verify function argument types, but for now we must verify their type
+        // manually
+        if (crewName.type != ExprValueType.STRING) {
+            val errorContext = PropertyValueMap().also {
+                it[Property.EXPECTED_ARGUMENT_TYPES] = "STRING"
+                it[Property.ACTUAL_ARGUMENT_TYPES] = crewName.type.name
+                it[Property.FUNCTION_NAME] = signature.name
+            }
+            throw EvaluationException("First argument to ${signature.name} was not a string",
+                ErrorCode.EVALUATOR_INCORRECT_TYPE_OF_ARGUMENTS_TO_PROCEDURE_CALL,
+                errorContext,
+                internal = false)
+        }
+
+        // Next we check if the given `crewName` is in the [EvaluationSession]'s global bindings. If not, we return 0.
+        val sessionGlobals = session.globals
+        val crewBindings = sessionGlobals[BindingName(crewName.stringValue(), BindingCase.INSENSITIVE)]
+            ?: return valueFactory.newInt(-1)
+
+        // Now that we've confirmed the given `crewName` is in the session's global bindings, we calculate and store
+        // the moon weight for each crewmate in the crew.
+        // In addition, we keep a running a tally of how many crewmates we do this for.
+        var numCalculated = 0
+        for (crewmateBinding in crewBindings) {
+            val crewmate = crewmateBinding.ionValue as IonStruct
+            val mass = crewmate["mass"] as IonDecimal
+            val moonWeight = (mass.decimalValue() * MOON_GRAVITATIONAL_CONSTANT).setScale(1, RoundingMode.HALF_UP)
+            crewmate.add("moonWeight", ion.newDecimal(moonWeight))
+
+            numCalculated++
+        }
+        return valueFactory.newInt(numCalculated)
+    }
+}
+
+/**
+ * Demonstrates the use of custom stored procedure [CalculateCrewMoonWeight] in PartiQL queries.
+ */
+class CustomProceduresExample(out: PrintStream) : Example(out) {
+    override fun run() {
+        /**
+         * To make custom stored procedures available to the PartiQL query being executed, they must be passed to
+         * [CompilerPipeline.Builder.addProcedure].
+         */
+        val pipeline = CompilerPipeline.build(ion) {
+            addProcedure(CalculateCrewMoonWeight(valueFactory))
+        }
+
+        // Here, we initialize the crews to be stored in our global session bindings
+        val initialCrews = Bindings.ofMap(
+            mapOf(
+                "crew1" to pipeline.valueFactory.newFromIonValue(
+                    ion.singleValue("""[ { name: "Neil",    mass: 80.5 }, 
+                                         { name: "Buzz",    mass: 72.3 },
+                                         { name: "Michael", mass: 89.9 } ]""")),
+                "crew2" to pipeline.valueFactory.newFromIonValue(
+                    ion.singleValue("""[ { name: "James", mass: 77.1 }, 
+                                         { name: "Spock", mass: 81.6 } ]"""))
+            )
+        )
+        val session = EvaluationSession.build { globals(initialCrews) }
+
+        val crew1BindingName = BindingName("crew1", BindingCase.INSENSITIVE)
+        val crew2BindingName = BindingName("crew2", BindingCase.INSENSITIVE)
+
+        out.println("Initial global session bindings:")
+        print("Crew 1:", "${session.globals[crew1BindingName]}")
+        print("Crew 2:", "${session.globals[crew2BindingName]}")
+
+        // We call our custom stored procedure using PartiQL's `EXEC` clause. Here we call our stored procedure
+        // 'calculate_crew_moon_weight' with the arg 'crew1', which outputs the number of crewmates we've calculated
+        // the moon weight for
+        val procedureCall = "EXEC calculate_crew_moon_weight 'crew1'"
+        val procedureCallOutput = pipeline.compile(procedureCall).eval(session)
+        print("Number of calculated moon weights:", "$procedureCallOutput")
+
+        out.println("Updated global session bindings:")
+        print("Crew 1:", "${session.globals[crew1BindingName]}")
+        print("Crew 2:", "${session.globals[crew2BindingName]}")
+    }
+}

examples/src/kotlin/org/partiql/examples/util/Main.kt

@@ -15,6 +15,7 @@ private val examples = mapOf(
         // Kotlin Examples
         CsvExprValueExample::class.java.simpleName to CsvExprValueExample(System.out),
         CustomFunctionsExample::class.java.simpleName to CustomFunctionsExample(System.out),
+        CustomProceduresExample::class.java.simpleName to CustomProceduresExample(System.out),
         EvaluationWithBindings::class.java.simpleName to EvaluationWithBindings(System.out),
         EvaluationWithLazyBindings::class.java.simpleName to EvaluationWithLazyBindings(System.out),
         ParserErrorExample::class.java.simpleName to ParserErrorExample(System.out),

examples/test/org/partiql/examples/CustomProceduresExampleTest.kt

@@ -0,0 +1,24 @@
+package org.partiql.examples
+
+import org.partiql.examples.util.Example
+import java.io.PrintStream
+
+class CustomProceduresExampleTest : BaseExampleTest() {
+    override fun example(out: PrintStream): Example = CustomProceduresExample(out)
+
+    override val expected = """
+        |Initial global session bindings:
+        |Crew 1:
+        |    [{'name': 'Neil', 'mass': 80.5}, {'name': 'Buzz', 'mass': 72.3}, {'name': 'Michael', 'mass': 89.9}]
+        |Crew 2:
+        |    [{'name': 'James', 'mass': 77.1}, {'name': 'Spock', 'mass': 81.6}]
+        |Number of calculated moon weights:
+        |    3
+        |Updated global session bindings:
+        |Crew 1:
+        |    [{'name': 'Neil', 'mass': 80.5, 'moonWeight': 13.3}, {'name': 'Buzz', 'mass': 72.3, 'moonWeight': 12.0}, {'name': 'Michael', 'mass': 89.9, 'moonWeight': 14.9}]
+        |Crew 2:
+        |    [{'name': 'James', 'mass': 77.1}, {'name': 'Spock', 'mass': 81.6}]
+        |
+    """.trimMargin()
+}

lang/src/org/partiql/lang/CompilerPipeline.kt

@@ -18,6 +18,7 @@ import com.amazon.ion.*
 import org.partiql.lang.ast.*
 import org.partiql.lang.eval.*
 import org.partiql.lang.eval.builtins.*
+import org.partiql.lang.eval.builtins.storedprocedure.StoredProcedure
 import org.partiql.lang.syntax.*
 
 /**
@@ -35,7 +36,13 @@ data class StepContext(
      * Includes built-in functions as well as custom functions added while the [CompilerPipeline]
      * was being built.
      */
-    val functions: @JvmSuppressWildcards Map<String, ExprFunction>
+    val functions: @JvmSuppressWildcards Map<String, ExprFunction>,
+
+    /**
+     * Returns a list of all stored procedures which are available for execution.
+     * Only includes the custom stored procedures added while the [CompilerPipeline] was being built.
+     */
+    val procedures: @JvmSuppressWildcards Map<String, StoredProcedure>
 )
 
 /**
@@ -61,6 +68,12 @@ interface CompilerPipeline  {
      */
     val functions: @JvmSuppressWildcards Map<String, ExprFunction>
 
+    /**
+     * Returns a list of all stored procedures which are available for execution.
+     * Only includes the custom stored procedures added while the [CompilerPipeline] was being built.
+     */
+    val procedures: @JvmSuppressWildcards Map<String, StoredProcedure>
+
     /** Compiles the specified PartiQL query using the configured parser. */
     fun compile(query: String): Expression
 
@@ -98,6 +111,7 @@ interface CompilerPipeline  {
         private var parser: Parser? = null
         private var compileOptions: CompileOptions? = null
         private val customFunctions: MutableMap<String, ExprFunction> = HashMap()
+        private val customProcedures: MutableMap<String, StoredProcedure> = HashMap()
         private val preProcessingSteps: MutableList<ProcessingStep> = ArrayList()
 
         /**
@@ -129,6 +143,13 @@ interface CompilerPipeline  {
          */
         fun addFunction(function: ExprFunction): Builder = this.apply { customFunctions[function.name] = function }
 
+        /**
+         * Add a custom stored procedure which will be callable by the compiled queries.
+         *
+         * Stored procedures added here will replace any built-in procedure with the same name.
+         */
+        fun addProcedure(procedure: StoredProcedure): Builder = this.apply { customProcedures[procedure.signature.name] = procedure }
+
         /** Adds a preprocessing step to be executed after parsing but before compilation. */
         fun addPreprocessingStep(step: ProcessingStep): Builder = this.apply { preProcessingSteps.add(step) }
 
@@ -145,6 +166,7 @@ interface CompilerPipeline  {
                 parser ?: SqlParser(valueFactory.ion),
                 compileOptions ?: CompileOptions.standard(),
                 allFunctions,
+                customProcedures,
                 preProcessingSteps)
         }
     }
@@ -155,17 +177,18 @@ private class CompilerPipelineImpl(
     private val parser: Parser,
     override val compileOptions: CompileOptions,
     override val functions: Map<String, ExprFunction>,
+    override val procedures: Map<String, StoredProcedure>,
     private val preProcessingSteps: List<ProcessingStep>
 ) : CompilerPipeline {
 
-    private val compiler = EvaluatingCompiler(valueFactory, functions, compileOptions)
+    private val compiler = EvaluatingCompiler(valueFactory, functions, procedures, compileOptions)
 
     override fun compile(query: String): Expression {
         return compile(parser.parseExprNode(query))
     }
 
     override fun compile(query: ExprNode): Expression {
-        val context = StepContext(valueFactory, compileOptions, functions)
+        val context = StepContext(valueFactory, compileOptions, functions, procedures)
 
         val preProcessedQuery = preProcessingSteps.fold(query) { currentExprNode, step ->
             step(currentExprNode, context)

lang/src/org/partiql/lang/errors/ErrorAndErrorContexts.kt

@@ -72,6 +72,7 @@ enum class Property(val propertyName: String, val propertyType: PropertyType) {
     LIKE_PATTERN("pattern", STRING_CLASS),
     LIKE_ESCAPE("escape_char", STRING_CLASS),
     FUNCTION_NAME("function_name", STRING_CLASS),
+    PROCEDURE_NAME("procedure_name", STRING_CLASS),
     EXPECTED_ARGUMENT_TYPES("expected_types", STRING_CLASS),
     ACTUAL_ARGUMENT_TYPES("actual_types", STRING_CLASS),
     FEATURE_NAME("FEATURE_NAME", STRING_CLASS),

lang/src/org/partiql/lang/errors/ErrorCode.kt

@@ -417,11 +417,24 @@ enum class ErrorCode(private val category: ErrorCategory,
                 "No such function: ${errorContext?.get(Property.FUNCTION_NAME)?.stringValue() ?: UNKNOWN} "
         },
 
+    EVALUATOR_NO_SUCH_PROCEDURE(
+        ErrorCategory.EVALUATOR,
+        LOCATION + setOf(Property.PROCEDURE_NAME),
+        ""){
+            override fun getErrorMessage(errorContext: PropertyValueMap?): String =
+                "No such stored procedure: ${errorContext?.get(Property.PROCEDURE_NAME)?.stringValue() ?: UNKNOWN} "
+        },
+
     EVALUATOR_INCORRECT_NUMBER_OF_ARGUMENTS_TO_FUNC_CALL(
         ErrorCategory.EVALUATOR,
         LOCATION + setOf(Property.EXPECTED_ARITY_MIN, Property.EXPECTED_ARITY_MAX),
         "Incorrect number of arguments to function call"),
 
+    EVALUATOR_INCORRECT_NUMBER_OF_ARGUMENTS_TO_PROCEDURE_CALL(
+        ErrorCategory.EVALUATOR,
+        LOCATION + setOf(Property.EXPECTED_ARITY_MIN, Property.EXPECTED_ARITY_MAX),
+        "Incorrect number of arguments to procedure call"),
+
     EVALUATOR_INCORRECT_TYPE_OF_ARGUMENTS_TO_FUNC_CALL(
         ErrorCategory.EVALUATOR,
         LOCATION + setOf(Property.EXPECTED_ARGUMENT_TYPES, Property.ACTUAL_ARGUMENT_TYPES, Property.FUNCTION_NAME),
@@ -432,6 +445,16 @@ enum class ErrorCode(private val category: ErrorCategory,
             "got: ${errorContext?.get(Property.ACTUAL_ARGUMENT_TYPES) ?: UNKNOWN}"
     },
 
+    EVALUATOR_INCORRECT_TYPE_OF_ARGUMENTS_TO_PROCEDURE_CALL(
+        ErrorCategory.EVALUATOR,
+        LOCATION + setOf(Property.EXPECTED_ARGUMENT_TYPES, Property.ACTUAL_ARGUMENT_TYPES, Property.FUNCTION_NAME),
+        "Incorrect type of arguments to procedure call") {
+        override fun getErrorMessage(errorContext: PropertyValueMap?): String =
+            "Invalid argument types for ${errorContext?.get(Property.FUNCTION_NAME) ?: UNKNOWN}, " +
+            "expected: ${errorContext?.get(Property.EXPECTED_ARGUMENT_TYPES) ?: UNKNOWN} " +
+            "got: ${errorContext?.get(Property.ACTUAL_ARGUMENT_TYPES) ?: UNKNOWN}"
+    },
+
     EVALUATOR_CONCAT_FAILED_DUE_TO_INCOMPATIBLE_TYPE(
         ErrorCategory.EVALUATOR,
         LOCATION + setOf(Property.ACTUAL_ARGUMENT_TYPES),

lang/src/org/partiql/lang/eval/EvaluatingCompiler.kt

@@ -21,6 +21,7 @@ import org.partiql.lang.ast.passes.*
 import org.partiql.lang.domains.PartiqlAst
 import org.partiql.lang.errors.*
 import org.partiql.lang.eval.binding.*
+import org.partiql.lang.eval.builtins.storedprocedure.StoredProcedure
 import org.partiql.lang.eval.like.PatternPart
 import org.partiql.lang.eval.like.executePattern
 import org.partiql.lang.eval.like.parsePattern
@@ -54,6 +55,7 @@ import kotlin.collections.*
 internal class EvaluatingCompiler(
     private val valueFactory: ExprValueFactory,
     private val functions: Map<String, ExprFunction>,
+    private val procedures: Map<String, StoredProcedure>,
     private val compileOptions: CompileOptions = CompileOptions.standard()
 ) {
     private val thunkFactory = ThunkFactory(compileOptions.thunkOptions)
@@ -1926,8 +1928,46 @@ internal class EvaluatingCompiler(
         }
     }
 
-    private fun compileExec(node: ExprNode): ThunkEnv {
-        TODO()
+    private fun compileExec(node: Exec): ThunkEnv {
+        val (procedureName, args, metas: MetaContainer) = node
+        val procedure = procedures[procedureName.name] ?: err(
+            "No such stored procedure: ${procedureName.name}",
+            ErrorCode.EVALUATOR_NO_SUCH_PROCEDURE,
+            errorContextFrom(metas).also {
+                it[Property.PROCEDURE_NAME] = procedureName.name
+            },
+            internal = false)
+
+        // Check arity
+        if (args.size !in procedure.signature.arity) {
+            val errorContext = errorContextFrom(metas).also {
+                it[Property.EXPECTED_ARITY_MIN] = procedure.signature.arity.first
+                it[Property.EXPECTED_ARITY_MAX] = procedure.signature.arity.last
+            }
+
+            val message = when {
+                procedure.signature.arity.first == 1 && procedure.signature.arity.last == 1 ->
+                    "${procedure.signature.name} takes a single argument, received: ${args.size}"
+                procedure.signature.arity.first == procedure.signature.arity.last ->
+                    "${procedure.signature.name} takes exactly ${procedure.signature.arity.first} arguments, received: ${args.size}"
+                else ->
+                    "${procedure.signature.name} takes between ${procedure.signature.arity.first} and " +
+                        "${procedure.signature.arity.last} arguments, received: ${args.size}"
+            }
+
+            throw EvaluationException(message,
+                ErrorCode.EVALUATOR_INCORRECT_NUMBER_OF_ARGUMENTS_TO_PROCEDURE_CALL,
+                errorContext,
+                internal = false)
+        }
+
+        // Compile the procedure's arguments
+        val argThunks = args.map { compileExprNode(it) }
+
+        return thunkFactory.thunkEnv(metas) { env ->
+            val procedureArgValues = argThunks.map { it(env) }
+            procedure.call(env.session, procedureArgValues)
+        }
     }
 
     /** A special wrapper for `UNPIVOT` values as a BAG. */