initial support for arrow serialization (finos#2258)

* initial support for arrow serialization
PrateekGarg-gs · Sep 20, 2023 · 83bc59e · 83bc59e
1 parent 3c0b09e
commit 83bc59e
Show file tree

Hide file tree

Showing 36 changed files with 1,544 additions and 29 deletions.
diff --git a/legend-engine-config/legend-engine-extensions-collection-execution/pom.xml b/legend-engine-config/legend-engine-extensions-collection-execution/pom.xml
@@ -158,6 +158,11 @@
             <groupId>org.finos.legend.engine</groupId>
             <artifactId>legend-engine-xt-flatdata-shared</artifactId>
         </dependency>
+
+        <dependency>
+            <groupId>org.finos.legend.engine</groupId>
+            <artifactId>legend-engine-xt-arrow-runtime</artifactId>
+        </dependency>
         <dependency>
             <groupId>org.finos.legend.engine</groupId>
             <artifactId>legend-engine-xt-flatdata-driver-bloomberg</artifactId>

diff --git a/...src/test/java/org/finos/legend/engine/extensions/collection/execution/TestExtensions.java b/...src/test/java/org/finos/legend/engine/extensions/collection/execution/TestExtensions.java
@@ -20,6 +20,7 @@
 import org.eclipse.collections.api.list.MutableList;
 import org.eclipse.collections.api.set.MutableSet;
 import org.eclipse.collections.impl.factory.Sets;
+import org.finos.legend.engine.external.format.arrow.ArrowRuntimeExtension;
 import org.finos.legend.engine.external.format.flatdata.FlatDataRuntimeExtension;
 import org.finos.legend.engine.external.format.flatdata.driver.spi.FlatDataDriverDescription;
 import org.finos.legend.engine.external.format.json.JsonSchemaRuntimeExtension;
@@ -98,7 +99,8 @@ protected MutableList<Class<? extends ExternalFormatRuntimeExtension>> expectedE
         return Lists.mutable.<Class<? extends ExternalFormatRuntimeExtension>>empty()
                 .with(FlatDataRuntimeExtension.class)
                 .with(JsonSchemaRuntimeExtension.class)
-                .with(XsdRuntimeExtension.class);
+                .with(XsdRuntimeExtension.class)
+                .with(ArrowRuntimeExtension.class);
     }
 
     @Test

diff --git a/legend-engine-config/legend-engine-extensions-collection-generation/pom.xml b/legend-engine-config/legend-engine-extensions-collection-generation/pom.xml
@@ -214,7 +214,12 @@
         </dependency>
         <!-- Flat Data -->
 
+        <!-- ARROW -->
 
+        <dependency>
+            <groupId>org.finos.legend.engine</groupId>
+            <artifactId>legend-engine-xt-arrow-pure</artifactId>
+        </dependency>
         <!-- JSON -->
         <dependency>
             <groupId>org.finos.legend.engine</groupId>

diff --git a/...rc/test/java/org/finos/legend/engine/extensions/collection/generation/TestExtensions.java b/...rc/test/java/org/finos/legend/engine/extensions/collection/generation/TestExtensions.java
@@ -414,6 +414,7 @@ protected Iterable<? extends Class<? extends PureCoreExtension>> getExpectedPlan
                 .with(XMLJavaBindingPureCoreExtension.class)
                 .with(ServicePureCoreExtension.class)
                 .with(RelationalJavaBindingPureCoreExtension.class)
+                .with(ArrowPureCoreExtension.class)
                 ;
     }
 
@@ -427,7 +428,8 @@ protected Iterable<? extends Class<? extends ExternalFormatExtension<?>>> getExp
                 .with(org.finos.legend.engine.external.format.protobuf.ProtobufFormatExtension.class)
                 .with(org.finos.legend.engine.query.graphQL.api.format.GraphQLFormatExtension.class)
                 .with(org.finos.legend.engine.query.graphQL.api.format.GraphQLSDLFormatExtension.class)
-                .with(org.finos.legend.engine.external.format.daml.DamlFormatExtension.class);
+                .with(org.finos.legend.engine.external.format.daml.DamlFormatExtension.class)
+                ;
     }
 
     protected Iterable<? extends Class<? extends FlatDataDriverDescription>> getExpectedFlatDataDriverDescriptionExtensions()
@@ -483,6 +485,7 @@ protected Iterable<String> getExpectedCodeRepositories()
                 .with("core_external_format_openapi")
                 .with("core_external_format_protobuf")
                 .with("core_external_format_xml")
+                .with("core_external_format_arrow")
                 .with("core_external_query_graphql")
                 .with("core_external_query_graphql_metamodel")
                 .with("core_external_query_sql_metamodel")

diff --git a/legend-engine-config/legend-engine-server/pom.xml b/legend-engine-config/legend-engine-server/pom.xml
@@ -503,6 +503,17 @@
             <artifactId>legend-engine-xt-flatdata-runtime</artifactId>
             <scope>runtime</scope>
         </dependency>
+        <dependency>
+            <groupId>org.finos.legend.engine</groupId>
+            <artifactId>legend-engine-xt-arrow-runtime</artifactId>
+            <scope>runtime</scope>
+        </dependency>
+        <dependency>
+            <groupId>org.finos.legend.engine</groupId>
+            <artifactId>legend-engine-xt-arrow-pure</artifactId>
+            <scope>runtime</scope>
+        </dependency>
+
         <dependency>
             <groupId>org.finos.legend.engine</groupId>
             <artifactId>legend-engine-xt-json-model</artifactId>

diff --git a/...ava/org/finos/legend/engine/external/shared/runtime/ExternalFormatExecutionExtension.java b/...ava/org/finos/legend/engine/external/shared/runtime/ExternalFormatExecutionExtension.java
@@ -33,6 +33,7 @@
 import org.finos.legend.engine.protocol.pure.v1.model.executionPlan.nodes.ExecutionNode;
 import org.finos.legend.engine.protocol.pure.v1.model.executionPlan.nodes.externalFormat.DataQualityExecutionNode;
 import org.finos.legend.engine.protocol.pure.v1.model.executionPlan.nodes.externalFormat.ExternalFormatExternalizeExecutionNode;
+import org.finos.legend.engine.protocol.pure.v1.model.executionPlan.nodes.externalFormat.ExternalFormatExternalizeTDSExecutionNode;
 import org.finos.legend.engine.protocol.pure.v1.model.executionPlan.nodes.externalFormat.ExternalFormatInternalizeExecutionNode;
 import org.finos.legend.engine.protocol.pure.v1.model.executionPlan.nodes.externalFormat.UrlStreamExecutionNode;
 import org.finos.legend.engine.shared.core.url.UrlFactory;
@@ -71,6 +72,11 @@ else if (executionNode instanceof ExternalFormatExternalizeExecutionNode)
             {
                 return executeExternalizeExecutionNode((ExternalFormatExternalizeExecutionNode) executionNode, pm, executionState);
             }
+
+            else if (executionNode instanceof ExternalFormatExternalizeTDSExecutionNode)
+            {
+                return executeExternalizeTDSExecutionNode((ExternalFormatExternalizeTDSExecutionNode) executionNode, pm, executionState);
+            }
             else
             {
                 return null;
@@ -104,6 +110,19 @@ private Result executeExternalizeExecutionNode(ExternalFormatExternalizeExecutio
         return extension.executeExternalizeExecutionNode(node, result, profiles, executionState);
     }
 
+    private Result executeExternalizeTDSExecutionNode(ExternalFormatExternalizeTDSExecutionNode node, MutableList<CommonProfile> profiles, ExecutionState executionState)
+    {
+        ExternalFormatRuntimeExtension extension = EXTENSIONS.get(node.contentType);
+        if (extension == null)
+        {
+            throw new IllegalStateException("No runtime extension for contentType " + node.contentType);
+        }
+
+        Result result = node.executionNodes().getAny().accept(new ExecutionNodeExecutor(profiles, executionState));
+        return extension.executeExternalizeTDSExecutionNode(node, result, profiles, executionState);
+    }
+
+
     private Result executeUrlStream(UrlStreamExecutionNode node, MutableList<CommonProfile> profiles, ExecutionState executionState)
     {
         try

diff --git a/.../java/org/finos/legend/engine/external/shared/runtime/ExternalFormatRuntimeExtension.java b/.../java/org/finos/legend/engine/external/shared/runtime/ExternalFormatRuntimeExtension.java
@@ -19,6 +19,7 @@
 import org.finos.legend.engine.plan.execution.result.Result;
 import org.finos.legend.engine.plan.execution.result.object.StreamingObjectResult;
 import org.finos.legend.engine.protocol.pure.v1.model.executionPlan.nodes.externalFormat.ExternalFormatExternalizeExecutionNode;
+import org.finos.legend.engine.protocol.pure.v1.model.executionPlan.nodes.externalFormat.ExternalFormatExternalizeTDSExecutionNode;
 import org.finos.legend.engine.protocol.pure.v1.model.executionPlan.nodes.externalFormat.ExternalFormatInternalizeExecutionNode;
 import org.pac4j.core.profile.CommonProfile;
 
@@ -38,4 +39,9 @@ default Result executeExternalizeExecutionNode(ExternalFormatExternalizeExecutio
     {
         throw new UnsupportedOperationException("Externalize not supported by format - " + node.contentType);
     }
+
+    default Result executeExternalizeTDSExecutionNode(ExternalFormatExternalizeTDSExecutionNode node, Result result, MutableList<CommonProfile> profiles, ExecutionState executionState)
+    {
+        throw new UnsupportedOperationException("Externalize TDS not supported by format - " + node.contentType);
+    }
 }
diff --git a/...ure/src/main/java/org/finos/legend/engine/protocol/pure/v1/CorePureProtocolExtension.java b/...ure/src/main/java/org/finos/legend/engine/protocol/pure/v1/CorePureProtocolExtension.java
@@ -29,6 +29,7 @@
 import org.finos.legend.engine.protocol.pure.v1.model.executionPlan.nodes.ExecutionNode;
 import org.finos.legend.engine.protocol.pure.v1.model.executionPlan.nodes.externalFormat.DataQualityExecutionNode;
 import org.finos.legend.engine.protocol.pure.v1.model.executionPlan.nodes.externalFormat.ExternalFormatExternalizeExecutionNode;
+import org.finos.legend.engine.protocol.pure.v1.model.executionPlan.nodes.externalFormat.ExternalFormatExternalizeTDSExecutionNode;
 import org.finos.legend.engine.protocol.pure.v1.model.executionPlan.nodes.externalFormat.ExternalFormatInternalizeExecutionNode;
 import org.finos.legend.engine.protocol.pure.v1.model.executionPlan.nodes.externalFormat.UrlStreamExecutionNode;
 import org.finos.legend.engine.protocol.pure.v1.model.executionPlan.nodes.externalFormat.VariableResolutionExecutionNode;
@@ -130,6 +131,8 @@ public List<Function0<List<ProtocolSubTypeInfo<?>>>> getExtraProtocolSubTypeInfo
                         .withSubtype(VariableResolutionExecutionNode.class, "varResolution")
                         .withSubtype(ExternalFormatInternalizeExecutionNode.class, "externalFormatInternalize")
                         .withSubtype(ExternalFormatExternalizeExecutionNode.class, "externalFormatExternalize")
+                        .withSubtype(ExternalFormatExternalizeTDSExecutionNode.class, "externalFormatExternalizeTDS")
+
                         .build(),
                 ProtocolSubTypeInfo.newBuilder(TestSuite.class)
                         .withSubtype(MappingTestSuite.class, "mappingTestSuite")

diff --git a/...1/model/executionPlan/nodes/externalFormat/ExternalFormatExternalizeTDSExecutionNode.java b/...1/model/executionPlan/nodes/externalFormat/ExternalFormatExternalizeTDSExecutionNode.java
@@ -0,0 +1,29 @@
+//  Copyright 2022 Goldman Sachs
+//
+//  Licensed under the Apache License, Version 2.0 (the "License");
+//  you may not use this file except in compliance with the License.
+//  You may obtain a copy of the License at
+//
+//       http://www.apache.org/licenses/LICENSE-2.0
+//
+//  Unless required by applicable law or agreed to in writing, software
+//  distributed under the License is distributed on an "AS IS" BASIS,
+//  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+//  See the License for the specific language governing permissions and
+//  limitations under the License.
+
+package org.finos.legend.engine.protocol.pure.v1.model.executionPlan.nodes.externalFormat;
+
+import org.finos.legend.engine.protocol.pure.v1.model.executionPlan.nodes.ExecutionNode;
+import org.finos.legend.engine.protocol.pure.v1.model.executionPlan.nodes.ExecutionNodeVisitor;
+
+public class ExternalFormatExternalizeTDSExecutionNode extends ExecutionNode
+{
+    public String contentType;
+
+    @Override
+    public <T> T accept(ExecutionNodeVisitor<T> executionNodeVisitor)
+    {
+        return executionNodeVisitor.visit(this);
+    }
+}
diff --git a/...led-core/src/main/resources/core/pure/binding/executionPlan/executionPlan_generation.pure b/...led-core/src/main/resources/core/pure/binding/executionPlan/executionPlan_generation.pure
@@ -144,29 +144,12 @@ function <<access.private>> meta::external::format::shared::executionPlan::exter
 {
    let parameters      = $fe.parametersValues->evaluateAndDeactivate();
    let children        = $parameters->at(0)->processValueSpecification($state, $extensions, $debug)->toOneMany();
-
-   let inScopeVars     = $state.inScopeVars;
-
-    let tree  =  #{TabularDataSet{rows}}#;
-
-   let inputType       = $parameters->at(0)->byPassValueSpecificationWrapper()->cast(@SimpleFunctionExpression).genericType.rawType;
-   let checked         = $inputType == meta::pure::dataQuality::Checked;
-
-   assert($tree.class == $inputType, | 'Input type \'' + $inputType->toOne()->elementToPath() + '\' and root tree type \'' + $tree.class->toOne()->elementToPath() + '\' for externalize does not match');
-
-   let bindingArg      = $fe.parametersValues->at(1)->byPassValueSpecificationWrapper();
-   let binding         = $bindingArg->cast(@InstanceValue).values->match([
-                                                                    b:meta::external::format::shared::binding::Binding[1] | $b,
-                                                                    s:String[1] | ^meta::external::format::shared::binding::Binding(name = 'generatedBinding', package = meta::external::format::shared::executionPlan, contentType = $s, modelUnit = ^meta::pure::model::unit::ModelUnit(packageableElementIncludes = extractPackageableElementFromTree($tree)))
-                                                                  ]);
-
-   ^ExternalFormatExternalizeExecutionNode
+   let contentType      = $fe.parametersValues->at(1)->byPassValueSpecificationWrapper()->cast(@InstanceValue).values->at(0)->cast(@String);
+   ^ExternalFormatExternalizeTDSExecutionNode
    (
       resultType      = ^ResultType(type=String),
       resultSizeRange = PureOne,
-      checked         = $checked,
-      binding         = $binding,
-      tree            = $tree,
+      contentType         = $contentType,
       executionNodes  = $children
    );
 }

diff --git a/...compiled-core/src/main/resources/core/pure/binding/executionPlan/executionPlan_print.pure b/...compiled-core/src/main/resources/core/pure/binding/executionPlan/executionPlan_print.pure
@@ -33,6 +33,14 @@ function meta::external::format::shared::executionPlan::toString::printPlanNodeT
          $node.implementation->printImplementation('implementation', $space+'  ', $extensions) +
          $space + ')\n'
       },
+      {node:ExternalFormatExternalizeTDSExecutionNode[1] |
+         'ExternalFormat_ExternalizeTDS\n' + 
+         $space + '(' + header($node, $space, $extensions) + '\n' +
+         $space + '  contentType = ' + $node.contentType + '\n' +
+         $node->childrenToString($space+'  ', $extensions) + '\n' +
+         $node.implementation->printImplementation('implementation', $space+'  ', $extensions) +
+         $space + ')\n'
+      },
       {node:ExternalFormatInternalizeExecutionNode[1] |
          'ExternalFormat_Internalize\n' +
          $space + '(' + header($node, $space, $extensions) + '\n' +

diff --git a/...ine-pure-code-compiled-core/src/main/resources/core/pure/binding/executionPlan/model.pure b/...ine-pure-code-compiled-core/src/main/resources/core/pure/binding/executionPlan/model.pure
@@ -34,6 +34,12 @@ Class meta::external::format::shared::executionPlan::ExternalFormatExternalizeEx
   config  : ExternalFormatExternalizeConfig[0..1];
 }
 
+Class meta::external::format::shared::executionPlan::ExternalFormatExternalizeTDSExecutionNode extends ExecutionNode
+{
+  contentType : String[1];
+  config  : ExternalFormatExternalizeConfig[0..1];
+}
+
 Class meta::external::format::shared::executionPlan::ExternalFormatInternalizeExecutionNode extends ExecutionNode
 {
   binding           : Binding[1];

diff --git a/...-core/src/main/resources/core/pure/binding/protocols/pure/vX_X_X/extension/extension.pure b/...-core/src/main/resources/core/pure/binding/protocols/pure/vX_X_X/extension/extension.pure
@@ -24,7 +24,9 @@ function meta::protocols::pure::vX_X_X::external::shared::format::serializerExte
          {mapping:Mapping[1], extensions:Extension[*] |
             [
                d:meta::external::format::shared::executionPlan::ExternalFormatExternalizeExecutionNode[1] | transformExternalFormatExternalizeExecutionNode($d, $mapping, $extensions),
-               s:meta::external::format::shared::executionPlan::ExternalFormatInternalizeExecutionNode[1] | transformExternalFormatInternalizeExecutionNode($s, $mapping, $extensions)
+               s:meta::external::format::shared::executionPlan::ExternalFormatInternalizeExecutionNode[1] | transformExternalFormatInternalizeExecutionNode($s, $mapping, $extensions),
+               t:meta::external::format::shared::executionPlan::ExternalFormatExternalizeTDSExecutionNode[1] | transformExternalFormatExternalizeTDSExecutionNode($t, $mapping, $extensions)
+
             ]
          }
    );

diff --git a/...core/src/main/resources/core/pure/binding/protocols/pure/vX_X_X/models/executionPlan.pure b/...core/src/main/resources/core/pure/binding/protocols/pure/vX_X_X/models/executionPlan.pure
@@ -25,6 +25,11 @@ Class meta::protocols::pure::vX_X_X::metamodel::external::shared::format::execut
    binding     : String[1];
 }
 
+Class meta::protocols::pure::vX_X_X::metamodel::external::shared::format::executionPlan::ExternalFormatExternalizeTDSExecutionNode extends ExecutionNode
+{
+   contentType : String[1];
+}
+
 Class meta::protocols::pure::vX_X_X::metamodel::external::shared::format::executionPlan::ExternalFormatInternalizeExecutionNode extends ExecutionNode
 {
    contentType       : String[1];

diff --git a/...e/src/main/resources/core/pure/binding/protocols/pure/vX_X_X/transfers/executionPlan.pure b/...e/src/main/resources/core/pure/binding/protocols/pure/vX_X_X/transfers/executionPlan.pure
@@ -29,6 +29,16 @@ function meta::protocols::pure::vX_X_X::transformation::fromPureGraph::external:
    );
 }
 
+function meta::protocols::pure::vX_X_X::transformation::fromPureGraph::external::shared::format::transformExternalFormatExternalizeTDSExecutionNode(node:meta::external::format::shared::executionPlan::ExternalFormatExternalizeTDSExecutionNode[1], mapping:meta::pure::mapping::Mapping[1], extensions:Extension[*]): ExecutionNode[1]
+{
+   ^meta::protocols::pure::vX_X_X::metamodel::external::shared::format::executionPlan::ExternalFormatExternalizeTDSExecutionNode(
+      _type             = 'externalFormatExternalizeTDS',
+      resultType        = $node.resultType->meta::protocols::pure::vX_X_X::transformation::fromPureGraph::executionPlan::transformResultType($mapping, $extensions),
+      resultSizeRange   = $node.resultSizeRange->map(s| $s->meta::protocols::pure::vX_X_X::transformation::fromPureGraph::domain::transformMultiplicity()),
+      contentType       = $node.contentType
+   );
+}
+
 function meta::protocols::pure::vX_X_X::transformation::fromPureGraph::external::shared::format::transformExternalFormatInternalizeExecutionNode(node:meta::external::format::shared::executionPlan::ExternalFormatInternalizeExecutionNode[1], mapping:meta::pure::mapping::Mapping[1], extensions:Extension[*]): ExecutionNode[1]
 {
    ^meta::protocols::pure::vX_X_X::metamodel::external::shared::format::executionPlan::ExternalFormatInternalizeExecutionNode(

diff --git a/.../legend-engine-pure-ide-light/src/main/java/org/finos/legend/engine/ide/PureIDELight.java b/.../legend-engine-pure-ide-light/src/main/java/org/finos/legend/engine/ide/PureIDELight.java
@@ -100,7 +100,8 @@ protected MutableList<RepositoryCodeStorage> buildRepositories(SourceLocationCon
                 .with(this.buildCore("legend-engine-xts-mongodb/legend-engine-xt-nonrelationalStore-mongodb-javaPlatformBinding-pure","nonrelational-mongodb-java-platform-binding"))
                 .with(this.buildCore("legend-engine-xts-service/legend-engine-language-pure-dsl-service-pure","service"))
                 .with(this.buildCore("legend-engine-xts-iceberg/legend-engine-xt-iceberg-pure","external-tableformat-iceberg"))
-               ;
+                .with(this.buildCore("legend-engine-xts-arrow/legend-engine-xt-arrow-pure", "external-format-arrow"))
+                ;
     }
 
     @Override