Support for arrow's uuid extension.

Author: cnuernber

.gitignore

@@ -35,4 +35,5 @@ java_test/test.arrow
 java_test/simulation*
 .idea
 .calva
-*.iml
+*.iml
+.dir-locals.el

build.clj

@@ -18,7 +18,7 @@
 (defn compile [_]
   (b/javac {:src-dirs ["java" "java_public_api"]
             :class-dir class-dir
-            :basis (b/create-basis {:project "deps.edn" :aliases [:dev-m1-mac]})
+            :basis (b/create-basis {:project "deps.edn" :aliases [:dev-mac-m1]})
             :javac-opts ["-source" "8" "-target" "8" "-Xlint:unchecked"]}))
 
 (defn jar [_]

deps.edn

@@ -150,7 +150,7 @@
     }
    :extra-paths ["neanderthal" "test"]}
 
-  :dev-m1-mac
+  :dev-mac-m1
   {:extra-deps
    {criterium/criterium {:mvn/version "0.4.5"}
     com.cognitect/transit-clj {:mvn/version "1.0.333"}

scripts/arrow_uuid.py

@@ -0,0 +1,12 @@
+import pyarrow as pa
+import uuid as uuid
+
+schema = pa.schema([pa.field('id', pa.uuid())])
+data = [uuid.UUID("8be643d6-0df7-4e5e-837c-f94170c87914").bytes,
+        uuid.UUID("24bc9cf4-e2e8-444f-bb2d-82394f33ff76").bytes,
+        uuid.UUID("e8149e1b-aef6-4671-b1b4-3b7a21eed92a").bytes]
+
+with pa.OSFile('test/data/uuid_ext.arrow', 'wb') as sink:
+    with pa.ipc.new_file(sink, schema=schema) as writer:
+        batch = pa.record_batch([data], schema=schema)
+        writer.write(batch)

scripts/run-tests-m1

@@ -1,5 +1,5 @@
 #!/bin/bash
 
 scripts/compile
-clojure -X:dev-m1-mac:codegen
-clojure -M:dev-m1-mac:test
+clojure -X:dev-mac-m1:codegen
+clojure -M:dev-mac-m1:test

src/tech/v3/libs/arrow.clj

@@ -119,7 +119,7 @@
            [org.apache.arrow.vector.types.pojo Field Schema ArrowType$Int
             ArrowType$Utf8 ArrowType$Timestamp ArrowType$Time DictionaryEncoding FieldType
             ArrowType$FloatingPoint ArrowType$Bool ArrowType$Date ArrowType$Duration
-            ArrowType$LargeUtf8 ArrowType$Null ArrowType$List ArrowType$Binary]
+            ArrowType$LargeUtf8 ArrowType$Null ArrowType$List ArrowType$Binary ArrowType$FixedSizeBinary]
            [org.apache.arrow.flatbuf CompressionType]
            [org.apache.arrow.vector.types MetadataVersion]
            [org.apache.arrow.vector.ipc WriteChannel]           
@@ -132,7 +132,7 @@
            [java.io OutputStream InputStream ByteArrayOutputStream ByteArrayInputStream]
            [java.nio ByteBuffer ByteOrder ShortBuffer IntBuffer LongBuffer DoubleBuffer
             FloatBuffer]
-           [java.util List ArrayList Map HashMap Map$Entry Iterator Set]
+           [java.util List ArrayList Map HashMap Map$Entry Iterator Set UUID]
            [java.util.concurrent ForkJoinTask]
            [java.time ZoneId]
            [java.nio.channels WritableByteChannel]
@@ -435,6 +435,9 @@ Dependent block frames are not supported!!")
   (datafy [this] {:id (.getId this)
                   :ordered? (.isOrdered this)
                   :index-type (datafy (.getIndexType this))})
+  ArrowType$FixedSizeBinary
+  (datafy [this] {:datatype :fixed-size-binary
+                  :byte-width (.getByteWidth this)})
   ArrowType$List
   (datafy [this]
     {:datatype :list}))
@@ -710,10 +713,7 @@ Dependent block frames are not supported!!")
                  (ft-fn (ArrowType$Utf8.) encoding)
                  ;;If no encoding is provided then just save the string as text
                  (ft-fn (ArrowType$Utf8.)))
-       :uuid (do
-               (when (== 1 (long (swap! uuid-warn-counter inc)))
-                 (log/warn "Columns of type UUID are converted to type text when serializing to Arrow"))
-               (ft-fn (ArrowType$Utf8.)))
+       :uuid (ft-fn (ArrowType$FixedSizeBinary. 16))
        :text (ft-fn (ArrowType$Utf8.))
        :encoded-text (ft-fn (ArrowType$Utf8.))))))
 
@@ -722,6 +722,9 @@ Dependent block frames are not supported!!")
   ^Field [dictionaries {strings-as-text? :strings-as-text?}
           col]
   (let [colmeta (meta col)
+        colmeta (if (identical? :uuid (get colmeta :datatype))
+                  (assoc colmeta ARROW_EXTENSION_NAME ARROW_UUID_NAME)
+                  colmeta)
         nullable? (boolean
                    (or (:nullable? colmeta)
                        (not (empty? (ds-proto/missing col)))))
@@ -1225,6 +1228,18 @@ Dependent block frames are not supported!!")
              (throw (Exception. "Numeric buffer missing concrete representation"))))])
       (case col-dt
         :boolean [(boolean-bytes cbuf)]
+        :uuid (let [data (byte-array (* 16 (dtype/ecount cbuf)))
+                    wbuf (-> (java.nio.ByteBuffer/wrap data)
+                             (.order java.nio.ByteOrder/BIG_ENDIAN))]
+                (reduce (fn [_ ^UUID v]
+                          (if v
+                            (do 
+                              (.putLong wbuf (.getMostSignificantBits v))
+                              (.putLong wbuf (.getLeastSignificantBits v)))
+                            (do
+                              (.putLong wbuf 0) (.putLong wbuf 0))))
+                        nil cbuf)
+                [(java.nio.ByteBuffer/wrap data)])
         :string (let [str-t (ds-base/ensure-column-string-table col)
                       indices (dtype-proto/->array-buffer (str-table/indices str-t))]
                   [(nio-buffer/as-nio-buffer indices)])
@@ -1633,8 +1648,45 @@ Dependent block frames are not supported!!")
          (field-metadata field)
          (node-buf->missing node validity-buf))))))
 
-(defmethod ^:private preparse-field :default
+(def ^{:private true
+       :tag String} ARROW_EXTENSION_NAME "ARROW:extension:name")
+(def ^{:private true
+       :tag String} ARROW_UUID_NAME "arrow.uuid")
+
+(defmethod ^:private preparse-field :fixed-size-binary
   [field ^Iterator node-iter ^Iterator buf-iter dict-map options]
+  (let [node (.next node-iter)
+        buffers [(.next buf-iter) (.next buf-iter)]        
+        n-elems (long (:n-elems node))
+        field-width (long (get-in field [:field-type :byte-width]))]
+    (fn parse-fixed-binary-field
+      [decompressor]
+      (let [[validity-buf data-buf] (decompressor buffers)
+            ^bytes data-ary (if (instance? NativeBuffer data-buf)
+                              (native-buffer/->jvm-array data-buf 0 (dtype/ecount data-buf))
+                              (dtype/->array data-buf))
+            fm (field-metadata field)]
+        (col-impl/new-column
+         (:name field)
+         (if (= ARROW_UUID_NAME (get fm ARROW_EXTENSION_NAME))
+           (let [longsdata (-> (java.nio.ByteBuffer/wrap data-ary)
+                               (.order (java.nio.ByteOrder/BIG_ENDIAN)))]
+             (println "is uuid")
+             (dtype/make-reader :uuid n-elems
+                                (let [lidx (* idx 16)]
+                                  (java.util.UUID. (.getLong longsdata lidx)
+                                                   (.getLong longsdata (+ lidx 8))))))
+           (let [ll (ArrayLists/toList data-ary)]
+             (println "is obj")
+             (dtype/make-reader :object n-elems
+                                (let [lidx (* idx field-width)]
+                                  (.subList ll lidx (+ lidx field-width))))))
+         fm
+         (node-buf->missing node validity-buf))))))
+
+
+(defmethod ^:private preparse-field :default
+  [field ^Iterator node-iter ^Iterator buf-iter dict-map options]  
   (assert (= 0 (count (:children field)))
           (format "Field %s cannot be parsed with default parser" field))
   (let [field-dtype (get-in field [:field-type :datatype])
@@ -2094,16 +2146,7 @@ Please use stream->dataset-seq.")))
     ;;datatypes
     (reduce
      (fn [ds col]
-       (cond
-         (= :uuid (dtype/elemwise-datatype col))
-         (let [missing (ds-proto/missing col)
-               metadata (meta col)]
-           (assoc ds (metadata :name)
-                  #:tech.v3.dataset{:data (mapv (comp #(Text. %) str) col)
-                                    :missing missing
-                                    :metadata metadata
-                                    :name (metadata :name)}))
-         (and (= :string (dtype/elemwise-datatype col))
+       (if (and (= :string (dtype/elemwise-datatype col))
               (not (:strings-as-text? options)))
          (if (and (nil? prev-ds)
                   (instance? StringTable (.data ^Column col)))
@@ -2131,7 +2174,6 @@ Please use stream->dataset-seq.")))
                                           :metadata (assoc metadata
                                                            ::previous-string-table prev-str-t)
                                           :name (metadata :name)})))))
-         :else
          ds))
      ds
      (ds-base/columns ds))))

test/data/uuid_ext.arrow

@@ -242,18 +242,26 @@
 
 
 (deftest uuid-test
+  (let [py-uuid (ds/->dataset "test/data/uuid_ext.arrow" {:key-fn keyword})]
+    
+    (is (= :uuid (dtype/elemwise-datatype (py-uuid :id))))
+    (is (= (mapv #(java.util.UUID/fromString %)
+                 ["8be643d6-0df7-4e5e-837c-f94170c87914"
+                  "24bc9cf4-e2e8-444f-bb2d-82394f33ff76"
+                  "e8149e1b-aef6-4671-b1b4-3b7a21eed92a"])
+           (py-uuid :id))))
   (try
     (let [uuid-ds (ds/->dataset "test/data/uuid.parquet"
                                 {:parser-fn {"uuids" :uuid}})
           _ (arrow/write-dataset-to-stream! uuid-ds "test-uuid.arrow")
           copying-ds (arrow/read-stream-dataset-copying "test-uuid.arrow")
           inplace-ds (arrow/read-stream-dataset-inplace "test-uuid.arrow")]
-      (is (= :text ((comp :datatype meta) (copying-ds "uuids"))))
-      (is (= :text ((comp :datatype meta) (inplace-ds "uuids"))))
+      (is (= :uuid ((comp :datatype meta) (copying-ds "uuids"))))
+      (is (= :uuid ((comp :datatype meta) (inplace-ds "uuids"))))
       (is (= (vec (copying-ds "uuids"))
              (vec (inplace-ds "uuids"))))
-      (is (= (mapv str (uuid-ds "uuids"))
-             (mapv str (copying-ds "uuids")))))
+      (is (= (vec (uuid-ds "uuids"))
+             (vec (copying-ds "uuids")))))
     (finally
       (.delete (java.io.File. "test-uuid.arrow")))))