ceph/src/arrow/go/arrow/cdata/test/test_cimport.go

   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() {}