ceph/src/arrow/go/arrow/ipc/cmd/arrow-json-integration-test/main.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 package main // import "github.com/apache/arrow/go/v6/arrow/ipc/cmd/arrow-json-integration-test"
  18
  19 import (
  20         "flag"
  21         "log"
  22         "os"
  23
  24         "github.com/apache/arrow/go/v6/arrow"
  25         "github.com/apache/arrow/go/v6/arrow/array"
  26         "github.com/apache/arrow/go/v6/arrow/arrio"
  27         "github.com/apache/arrow/go/v6/arrow/internal/arrjson"
  28         "github.com/apache/arrow/go/v6/arrow/internal/testing/types"
  29         "github.com/apache/arrow/go/v6/arrow/ipc"
  30         "golang.org/x/xerrors"
  31 )
  32
  33 func main() {
  34         log.SetPrefix("arrow-json: ")
  35         log.SetFlags(0)
  36
  37         var (
  38                 arrowPath = flag.String("arrow", "", "path to ARROW file")
  39                 jsonPath  = flag.String("json", "", "path to JSON file")
  40                 mode      = flag.String("mode", "VALIDATE", "mode of integration testing tool (ARROW_TO_JSON, JSON_TO_ARROW, VALIDATE)")
  41                 verbose   = flag.Bool("verbose", true, "enable/disable verbose mode")
  42         )
  43
  44         flag.Parse()
  45
  46         err := runCommand(*jsonPath, *arrowPath, *mode, *verbose)
  47         if err != nil {
  48                 log.Fatal(err)
  49         }
  50 }
  51
  52 func runCommand(jsonName, arrowName, mode string, verbose bool) error {
  53         arrow.RegisterExtensionType(types.NewUUIDType())
  54
  55         if jsonName == "" {
  56                 return xerrors.Errorf("must specify json file name")
  57         }
  58
  59         if arrowName == "" {
  60                 return xerrors.Errorf("must specify arrow file name")
  61         }
  62
  63         switch mode {
  64         case "ARROW_TO_JSON":
  65                 return cnvToJSON(arrowName, jsonName, verbose)
  66         case "JSON_TO_ARROW":
  67                 return cnvToARROW(arrowName, jsonName, verbose)
  68         case "VALIDATE":
  69                 return validate(arrowName, jsonName, verbose)
  70         default:
  71                 return xerrors.Errorf("unknown command %q", mode)
  72         }
  73
  74         return nil
  75 }
  76
  77 func cnvToJSON(arrowName, jsonName string, verbose bool) error {
  78         r, err := os.Open(arrowName)
  79         if err != nil {
  80                 return xerrors.Errorf("could not open ARROW file %q: %w", arrowName, err)
  81         }
  82         defer r.Close()
  83
  84         w, err := os.Create(jsonName)
  85         if err != nil {
  86                 return xerrors.Errorf("could not create JSON file %q: %w", jsonName, err)
  87         }
  88         defer w.Close()
  89
  90         rr, err := ipc.NewFileReader(r)
  91         if err != nil {
  92                 return xerrors.Errorf("could not open ARROW file reader from file %q: %w", arrowName, err)
  93         }
  94         defer rr.Close()
  95
  96         if verbose {
  97                 log.Printf("found schema:\n%v\n", rr.Schema())
  98         }
  99
 100         ww, err := arrjson.NewWriter(w, rr.Schema())
 101         if err != nil {
 102                 return xerrors.Errorf("could not create JSON encoder: %w", err)
 103         }
 104         defer ww.Close()
 105
 106         n, err := arrio.Copy(ww, rr)
 107         if err != nil {
 108                 return xerrors.Errorf("could not convert ARROW file reader data to JSON data: %w", err)
 109         }
 110
 111         if got, want := n, int64(rr.NumRecords()); got != want {
 112                 return xerrors.Errorf("invalid number of records copied (got=%d, want=%d", got, want)
 113         }
 114
 115         err = ww.Close()
 116         if err != nil {
 117                 return xerrors.Errorf("could not close JSON encoder %q: %w", jsonName, err)
 118         }
 119
 120         err = w.Close()
 121         if err != nil {
 122                 return xerrors.Errorf("could not close JSON file %q: %w", jsonName, err)
 123         }
 124
 125         return nil
 126 }
 127
 128 func cnvToARROW(arrowName, jsonName string, verbose bool) error {
 129         r, err := os.Open(jsonName)
 130         if err != nil {
 131                 return xerrors.Errorf("could not open JSON file %q: %w", jsonName, err)
 132         }
 133         defer r.Close()
 134
 135         w, err := os.Create(arrowName)
 136         if err != nil {
 137                 return xerrors.Errorf("could not create ARROW file %q: %w", arrowName, err)
 138         }
 139         defer w.Close()
 140
 141         rr, err := arrjson.NewReader(r)
 142         if err != nil {
 143                 return xerrors.Errorf("could not open JSON file reader from file %q: %w", jsonName, err)
 144         }
 145
 146         if verbose {
 147                 log.Printf("found schema:\n%v\n", rr.Schema())
 148         }
 149
 150         ww, err := ipc.NewFileWriter(w, ipc.WithSchema(rr.Schema()))
 151         if err != nil {
 152                 return xerrors.Errorf("could not create ARROW file writer: %w", err)
 153         }
 154         defer ww.Close()
 155
 156         n, err := arrio.Copy(ww, rr)
 157         if err != nil {
 158                 return xerrors.Errorf("could not convert JSON data to ARROW data: %w", err)
 159         }
 160
 161         if got, want := n, int64(rr.NumRecords()); got != want {
 162                 return xerrors.Errorf("invalid number of records copied (got=%d, want=%d", got, want)
 163         }
 164
 165         err = ww.Close()
 166         if err != nil {
 167                 return xerrors.Errorf("could not close ARROW file writer %q: %w", arrowName, err)
 168         }
 169
 170         err = w.Close()
 171         if err != nil {
 172                 return xerrors.Errorf("could not close ARROW file %q: %w", arrowName, err)
 173         }
 174
 175         return nil
 176 }
 177
 178 func validate(arrowName, jsonName string, verbose bool) error {
 179         jr, err := os.Open(jsonName)
 180         if err != nil {
 181                 return xerrors.Errorf("could not open JSON file %q: %w", jsonName, err)
 182         }
 183         defer jr.Close()
 184
 185         jrr, err := arrjson.NewReader(jr)
 186         if err != nil {
 187                 return xerrors.Errorf("could not open JSON file reader from file %q: %w", jsonName, err)
 188         }
 189
 190         ar, err := os.Open(arrowName)
 191         if err != nil {
 192                 return xerrors.Errorf("could not open ARROW file %q: %w", arrowName, err)
 193         }
 194         defer ar.Close()
 195
 196         arr, err := ipc.NewFileReader(ar)
 197         if err != nil {
 198                 return xerrors.Errorf("could not open ARROW file reader from file %q: %w", arrowName, err)
 199         }
 200         defer arr.Close()
 201
 202         if !arr.Schema().Equal(jrr.Schema()) {
 203                 if verbose {
 204                         log.Printf("JSON schema:\n%v\nArrow schema:\n%v\n", arr.Schema(), jrr.Schema())
 205                 }
 206                 return xerrors.Errorf("schemas did not match")
 207         }
 208
 209         for i := 0; i < arr.NumRecords(); i++ {
 210                 arec, err := arr.Read()
 211                 if err != nil {
 212                         return xerrors.Errorf("could not read record %d from ARROW file: %w", i, err)
 213                 }
 214                 jrec, err := jrr.Read()
 215                 if err != nil {
 216                         return xerrors.Errorf("could not read record %d from JSON file: %w", i, err)
 217                 }
 218                 if !array.RecordApproxEqual(jrec, arec) {
 219                         return xerrors.Errorf("record batch %d did not match\nJSON:\n%v\nARROW:\n%v\n",
 220                                 i, jrec, arec,
 221                         )
 222                 }
 223         }
 224
 225         if jn, an := jrr.NumRecords(), arr.NumRecords(); jn != an {
 226                 return xerrors.Errorf("different number of record batches: %d (JSON) vs %d (Arrow)", jn, an)
 227         }
 228
 229         return nil
 230 }