]>
Commit | Line | Data |
---|---|---|
1d09f67e TL |
1 | // Licensed to the Apache Software Foundation (ASF) under one |
2 | // or more contributor license agreements. See the NOTICE file | |
3 | // distributed with this work for additional information | |
4 | // regarding copyright ownership. The ASF licenses this file | |
5 | // to you under the Apache License, Version 2.0 (the | |
6 | // "License"); you may not use this file except in compliance | |
7 | // with the License. You may obtain a copy of the License at | |
8 | // | |
9 | // http://www.apache.org/licenses/LICENSE-2.0 | |
10 | // | |
11 | // Unless required by applicable law or agreed to in writing, software | |
12 | // distributed under the License is distributed on an "AS IS" BASIS, | |
13 | // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. | |
14 | // See the License for the specific language governing permissions and | |
15 | // limitations under the License. | |
16 | ||
17 | // +build cdata_test | |
18 | ||
19 | package main | |
20 | ||
21 | import ( | |
22 | "fmt" | |
23 | "runtime" | |
24 | ||
25 | "github.com/apache/arrow/go/v6/arrow" | |
26 | "github.com/apache/arrow/go/v6/arrow/array" | |
27 | "github.com/apache/arrow/go/v6/arrow/cdata" | |
28 | "github.com/apache/arrow/go/v6/arrow/memory" | |
29 | ) | |
30 | ||
31 | // #include <stdint.h> | |
32 | import "C" | |
33 | ||
34 | var alloc = memory.NewCheckedAllocator(memory.NewGoAllocator()) | |
35 | ||
36 | //export totalAllocated | |
37 | func totalAllocated() int64 { | |
38 | return int64(alloc.CurrentAlloc()) | |
39 | } | |
40 | ||
41 | //export runGC | |
42 | func runGC() { | |
43 | runtime.GC() | |
44 | } | |
45 | ||
46 | //export importSchema | |
47 | func importSchema(ptr uintptr) { | |
48 | schema, err := cdata.ImportCArrowSchema(cdata.SchemaFromPtr(ptr)) | |
49 | if err != nil { | |
50 | panic(err) | |
51 | } | |
52 | ||
53 | expectedMetadata := arrow.NewMetadata([]string{"key1"}, []string{"value1"}) | |
54 | expectedSchema := arrow.NewSchema([]arrow.Field{{Name: "ints", Type: arrow.ListOf(arrow.PrimitiveTypes.Int32), Nullable: true}}, &expectedMetadata) | |
55 | if !schema.Equal(expectedSchema) { | |
56 | panic(fmt.Sprintf("schema didn't match: expected %s, got %s", expectedSchema, schema)) | |
57 | } | |
58 | if !schema.Metadata().Equal(expectedMetadata) { | |
59 | panic(fmt.Sprintf("metadata didn't match: expected %s, got %s", expectedMetadata, schema.Metadata())) | |
60 | } | |
61 | ||
62 | fmt.Println("schema matches! Huzzah!") | |
63 | } | |
64 | ||
65 | //export importRecordBatch | |
66 | func importRecordBatch(scptr, rbptr uintptr) { | |
67 | sc := cdata.SchemaFromPtr(scptr) | |
68 | rb := cdata.ArrayFromPtr(rbptr) | |
69 | ||
70 | rec, err := cdata.ImportCRecordBatch(rb, sc) | |
71 | if err != nil { | |
72 | panic(err) | |
73 | } | |
74 | defer rec.Release() | |
75 | ||
76 | expectedMetadata := arrow.NewMetadata([]string{"key1"}, []string{"value1"}) | |
77 | expectedSchema := arrow.NewSchema([]arrow.Field{{Name: "ints", Type: arrow.ListOf(arrow.PrimitiveTypes.Int32), Nullable: true}}, &expectedMetadata) | |
78 | ||
79 | bldr := array.NewRecordBuilder(alloc, expectedSchema) | |
80 | defer bldr.Release() | |
81 | ||
82 | lb := bldr.Field(0).(*array.ListBuilder) | |
83 | vb := lb.ValueBuilder().(*array.Int32Builder) | |
84 | ||
85 | // [[[1], [], None [2, 42]]] | |
86 | lb.Append(true) | |
87 | vb.Append(int32(1)) | |
88 | ||
89 | lb.Append(true) | |
90 | lb.Append(false) | |
91 | ||
92 | lb.Append(true) | |
93 | vb.AppendValues([]int32{2, 42}, nil) | |
94 | ||
95 | expectedRec := bldr.NewRecord() | |
96 | defer expectedRec.Release() | |
97 | ||
98 | if !array.RecordEqual(expectedRec, rec) { | |
99 | panic(fmt.Sprintf("records didn't match: expected %s\n got %s", expectedRec, rec)) | |
100 | } | |
101 | ||
102 | fmt.Println("record batch matches huzzah!") | |
103 | } | |
104 | ||
105 | func makeSchema() *arrow.Schema { | |
106 | meta := arrow.NewMetadata([]string{"key1"}, []string{"value1"}) | |
107 | return arrow.NewSchema([]arrow.Field{ | |
108 | {Name: "ints", Type: arrow.ListOf(arrow.PrimitiveTypes.Int32), Nullable: true}, | |
109 | }, &meta) | |
110 | } | |
111 | ||
112 | func makeBatch() array.Record { | |
113 | bldr := array.NewRecordBuilder(alloc, makeSchema()) | |
114 | defer bldr.Release() | |
115 | ||
116 | fbldr := bldr.Field(0).(*array.ListBuilder) | |
117 | valbldr := fbldr.ValueBuilder().(*array.Int32Builder) | |
118 | ||
119 | fbldr.Append(true) | |
120 | valbldr.Append(1) | |
121 | ||
122 | fbldr.Append(true) | |
123 | fbldr.AppendNull() | |
124 | fbldr.Append(true) | |
125 | valbldr.Append(2) | |
126 | valbldr.Append(42) | |
127 | ||
128 | return bldr.NewRecord() | |
129 | } | |
130 | ||
131 | //export exportSchema | |
132 | func exportSchema(schema uintptr) { | |
133 | cdata.ExportArrowSchema(makeSchema(), cdata.SchemaFromPtr(schema)) | |
134 | } | |
135 | ||
136 | //export exportRecordBatch | |
137 | func exportRecordBatch(schema, record uintptr) { | |
138 | batch := makeBatch() | |
139 | defer batch.Release() | |
140 | ||
141 | cdata.ExportArrowRecordBatch(batch, cdata.ArrayFromPtr(record), cdata.SchemaFromPtr(schema)) | |
142 | } | |
143 | ||
144 | //export importThenExportSchema | |
145 | func importThenExportSchema(input, output uintptr) { | |
146 | schema, err := cdata.ImportCArrowSchema(cdata.SchemaFromPtr(input)) | |
147 | if err != nil { | |
148 | panic(err) | |
149 | } | |
150 | ||
151 | cdata.ExportArrowSchema(schema, cdata.SchemaFromPtr(output)) | |
152 | } | |
153 | ||
154 | //export importThenExportRecord | |
155 | func importThenExportRecord(schemaIn, arrIn uintptr, schemaOut, arrOut uintptr) { | |
156 | rec, err := cdata.ImportCRecordBatch(cdata.ArrayFromPtr(arrIn), cdata.SchemaFromPtr(schemaIn)) | |
157 | if err != nil { | |
158 | panic(err) | |
159 | } | |
160 | ||
161 | defer rec.Release() | |
162 | cdata.ExportArrowRecordBatch(rec, cdata.ArrayFromPtr(arrOut), cdata.SchemaFromPtr(schemaOut)) | |
163 | } | |
164 | ||
165 | func main() {} |