Element encoding for doubles

ecmwf · Nov 26, 2024 · abba3a6 · abba3a6
1 parent f4279fc
commit abba3a6
Show file tree

Hide file tree

Showing 3 changed files with 64 additions and 1 deletion.
diff --git a/src/accessor/grib_accessor_class_element.cc b/src/accessor/grib_accessor_class_element.cc
@@ -123,6 +123,53 @@ int grib_accessor_element_t::pack_long(const long* val, size_t* len)
     return ret;
 }
 
+int grib_accessor_element_t::pack_double(const double* v, size_t* len)
+{
+    int ret               = 0;
+    size_t size           = 0;
+    double* ar            = NULL;
+    const grib_context* c = context_;
+    grib_handle* hand     = grib_handle_of_accessor(this);
+    long index            = element_;
+
+    if (*len < 1) {
+        ret = GRIB_ARRAY_TOO_SMALL;
+        return ret;
+    }
+
+    if ((ret = grib_get_size(hand, array_, &size)) != GRIB_SUCCESS)
+        return ret;
+
+    ar = (double*)grib_context_malloc_clear(c, size * sizeof(double));
+    if (!ar) {
+        grib_context_log(c, GRIB_LOG_ERROR, "Error allocating %zu bytes", size * sizeof(double));
+        return GRIB_OUT_OF_MEMORY;
+    }
+
+    if ((ret = grib_get_double_array_internal(hand, array_, ar, &size)) != GRIB_SUCCESS)
+        return ret;
+
+    // An index of -x means the xth item from the end of the list, so ar[-1] means the last item in ar
+    if (index < 0) {
+        index = size + index;
+    }
+
+    if ((ret = check_element_index(__func__, array_, index, size)) != GRIB_SUCCESS) {
+        goto the_end;
+    }
+
+    Assert(index >= 0);
+    Assert(index < size);
+    ar[index] = *v;
+
+    if ((ret = grib_set_double_array_internal(hand, array_, ar, size)) != GRIB_SUCCESS)
+        goto the_end;
+
+the_end:
+    grib_context_free(c, ar);
+    return ret;
+}
+
 int grib_accessor_element_t::unpack_double(double* val, size_t* len)
 {
     int ret                 = 0;

diff --git a/src/accessor/grib_accessor_class_element.h b/src/accessor/grib_accessor_class_element.h
@@ -21,6 +21,7 @@ class grib_accessor_element_t : public grib_accessor_long_t
     int pack_long(const long* val, size_t* len) override;
     int unpack_double(double* val, size_t* len) override;
     int unpack_long(long* val, size_t* len) override;
+    int pack_double(const double* val, size_t* len) override;
     void init(const long, grib_arguments*) override;
 
 private:

diff --git a/tests/grib_element.sh b/tests/grib_element.sh
@@ -11,6 +11,7 @@
 . ./include.ctest.sh
 
 label="grib_element_test"
+tempGrib=temp.${label}.grib
 tempRef=temp.${label}.ref
 tempText=temp.${label}.txt
 tempFilt=temp.${label}.filt
@@ -72,6 +73,20 @@ status=$?
 set -e
 [ $status -ne 0 ]
 
+# Encode a single double element in an array
+input=$data_dir/sample.grib2
+cat > $tempFilt <<EOF
+    meta elemN element(values, -1);
+    assert( elemN < 301 );
+    print "before: [elemN:d]";
+    set elemN = 311.0;  # value has to be a double
+    print "after:  [elemN:d]";
+    assert( elemN > 310 );
+    write;
+EOF
+${tools_dir}/grib_filter -o $tempGrib $tempFilt $input
+${tools_dir}/grib_get -p avg $input $tempGrib
+
 
 # Clean up
-rm -f $tempRef $tempText $tempFilt
+rm -f $tempRef $tempText $tempFilt $tempGrib