basic read-write test

Author: myrrc

vortex-ffi/cbindgen.toml

@@ -8,6 +8,7 @@ style = "type"
 header = """
 // SPDX-License-Identifier: Apache-2.0
 // SPDX-FileCopyrightText: Copyright the Vortex contributors
+#include <stdint.h>
 
 //
 // THIS FILE IS AUTO-GENERATED, DO NOT MAKE EDITS DIRECTLY

vortex-ffi/cinclude/vortex.h

@@ -1,5 +1,6 @@
 // SPDX-License-Identifier: Apache-2.0
 // SPDX-FileCopyrightText: Copyright the Vortex contributors
+#include <stdint.h>
 
 //
 // THIS FILE IS AUTO-GENERATED, DO NOT MAKE EDITS DIRECTLY
@@ -1345,6 +1346,7 @@ void vx_struct_column_builder_free(vx_struct_column_builder *ptr);
 /**
  * Create a new column-wise struct array builder with given validity and a
  * capacity hint. validity can't be NULL.
+ * Capacity hint is for the number of columns.
  * If you don't know capacity, pass 0.
  * if validity is NULL, returns NULL.
  */

vortex-ffi/src/array.rs

@@ -192,6 +192,9 @@ pub unsafe extern "C-unwind" fn vx_array_dtype(array: *const vx_array) -> *const
     vx_dtype::new_ref(vx_array::as_ref(array).dtype())
 }
 
+// Return an owned field for array at index.
+// Returns NULL and sets error_out if index is out of bounds or array doesn't
+// have dtype DTYPE_STRUCT.
 #[unsafe(no_mangle)]
 pub unsafe extern "C-unwind" fn vx_array_get_field(
     array: *const vx_array,

vortex-ffi/src/struct_array.rs

@@ -28,6 +28,7 @@ crate::box_wrapper!(StructBuilder, vx_struct_column_builder);
 
 /// Create a new column-wise struct array builder with given validity and a
 /// capacity hint. validity can't be NULL.
+/// Capacity hint is for the number of columns.
 /// If you don't know capacity, pass 0.
 /// if validity is NULL, returns NULL.
 #[unsafe(no_mangle)]

vortex-ffi/test/scan.cpp

@@ -18,39 +18,78 @@ struct TempPath : fs::path {
     }
 };
 
-constexpr size_t SAMPLE_ROWS = 0;
-[[nodiscard]] TempPath write_sample(vx_session *session, fs::path &&path) {
-    REQUIRE(path.is_absolute());
+// StructArray { age: u8, height: u16? }
+[[nodiscard]] const vx_dtype *sample_dtype() {
     vx_struct_fields_builder *builder = vx_struct_fields_builder_new();
 
-    {
-        constexpr auto col1 = "col1"sv;
-        const vx_string *col1_name = vx_string_new(col1.data(), col1.size());
-        const vx_dtype *col1_dtype = vx_dtype_new_primitive(PTYPE_U8, false);
-        vx_struct_fields_builder_add_field(builder, col1_name, col1_dtype);
+    constexpr auto age = "age"sv;
+    const vx_string *age_name = vx_string_new(age.data(), age.size());
+    const vx_dtype *age_type = vx_dtype_new_primitive(PTYPE_U8, false);
+    vx_struct_fields_builder_add_field(builder, age_name, age_type);
+
+    constexpr auto height = "height"sv;
+    const vx_string *height_name = vx_string_new(height.data(), height.size());
+    const vx_dtype *height_type = vx_dtype_new_primitive(PTYPE_U16, true);
+    vx_struct_fields_builder_add_field(builder, height_name, height_type);
+
+    vx_struct_fields *fields = vx_struct_fields_builder_finalize(builder);
+    return vx_dtype_new_struct(fields, false);
+}
+
+constexpr size_t SAMPLE_ROWS = 100;
+[[nodiscard]] const vx_array *sample_array() {
+    vx_validity validity = {};
+    validity.type = VX_VALIDITY_NON_NULLABLE;
+
+    vx_struct_column_builder *builder = vx_struct_column_builder_new(&validity, SAMPLE_ROWS);
+
+    vx_error *error = nullptr;
+
+    std::vector<uint8_t> age_buffer;
+    for (uint8_t age = 0; age < SAMPLE_ROWS; ++age) {
+        age_buffer.push_back(age);
     }
-    {
-        constexpr auto col2 = "col2"sv;
-        const vx_string *col2_name = vx_string_new(col2.data(), col2.size());
-        const vx_dtype *col2_dtype = vx_dtype_new_utf8(true);
-        vx_struct_fields_builder_add_field(builder, col2_name, col2_dtype);
+    const vx_array *age_array =
+        vx_array_new_primitive(PTYPE_U8, age_buffer.data(), age_buffer.size(), &validity, &error);
+    require_no_error(error);
+
+    vx_struct_column_builder_add_field(builder, "age", age_array, &error);
+    require_no_error(error);
+    vx_array_free(age_array);
+
+    std::vector<uint16_t> height_buffer;
+    for (uint16_t height = 0; height < SAMPLE_ROWS; ++height) {
+        height_buffer.push_back(1 + rand() % (height + 1));
     }
+    validity.type = VX_VALIDITY_ALL_VALID;
+    const vx_array *height_array =
+        vx_array_new_primitive(PTYPE_U16, height_buffer.data(), height_buffer.size(), &validity, &error);
+    require_no_error(error);
 
-    vx_struct_fields *fields = vx_struct_fields_builder_finalize(builder);
-    const vx_dtype *file_dtype = vx_dtype_new_struct(fields, false);
+    vx_struct_column_builder_add_field(builder, "height", height_array, &error);
+    require_no_error(error);
+    vx_array_free(height_array);
+
+    const vx_array *array = vx_struct_column_builder_finalize(builder, &error);
+    require_no_error(error);
+    return array;
+}
+
+[[nodiscard]] TempPath write_sample(vx_session *session, fs::path &&path) {
+    REQUIRE(path.is_absolute());
+
+    const vx_dtype *dtype = sample_dtype();
 
     vx_error *error = nullptr;
-    vx_array_sink *sink = vx_array_sink_open_file(session, path.c_str(), file_dtype, &error);
+    vx_array_sink *sink = vx_array_sink_open_file(session, path.c_str(), dtype, &error);
     REQUIRE(sink != nullptr);
     require_no_error(error);
-    vx_dtype_free(file_dtype);
+    vx_dtype_free(dtype);
 
-    for (size_t i = 0; i < SAMPLE_ROWS; ++i) {
-        //const vx_array* array = vx_array_new_primitive();
-        //vx_array_sink_push(sink, array, &error);
-        //require_no_error(error);
-        //vx_array_free(array);
-    }
+    const vx_array *array = sample_array();
+    vx_array_sink_push(sink, array, &error);
+    require_no_error(error);
+    vx_array_free(array);
 
     vx_array_sink_close(sink, &error);
     require_no_error(error);
@@ -110,6 +149,12 @@ TEST_CASE("Creating datasources", "[datasource]") {
     vx_session_free(session);
 }
 
+TEST_CASE("Write file", "[datasource]") {
+    vx_session *session = vx_session_new();
+    TempPath path = write_sample(session, fs::current_path() / "write-file.vortex");
+    vx_session_free(session);
+}
+
 TEST_CASE("Write file and read back types", "[datasource]") {
     vx_session *session = vx_session_new();
     TempPath path = write_sample(session, fs::current_path() / "write-read-types.vortex");
@@ -142,43 +187,143 @@ TEST_CASE("Write file and read back types", "[datasource]") {
     REQUIRE(vx_dtype_get_variant(col1_dtype) == DTYPE_PRIMITIVE);
     REQUIRE(vx_dtype_primitive_ptype(col1_dtype) == PTYPE_U8);
     REQUIRE_FALSE(vx_dtype_is_nullable(col1_dtype));
-    REQUIRE(to_string_view(col1_name) == "col1");
+    REQUIRE(to_string_view(col1_name) == "age");
     vx_dtype_free(col1_dtype);
 
     const vx_dtype *col2_dtype = vx_struct_fields_field_dtype(fields, 1);
     const vx_string *col2_name = vx_struct_fields_field_name(fields, 1);
 
-    REQUIRE(vx_dtype_get_variant(col2_dtype) == DTYPE_UTF8);
+    REQUIRE(vx_dtype_get_variant(col2_dtype) == DTYPE_PRIMITIVE);
+    REQUIRE(vx_dtype_primitive_ptype(col2_dtype) == PTYPE_U16);
     REQUIRE(vx_dtype_is_nullable(col2_dtype));
-    REQUIRE(to_string_view(col2_name) == "col2");
+    REQUIRE(to_string_view(col2_name) == "height");
     vx_dtype_free(col2_dtype);
 
     vx_data_source_free(ds);
 }
 
-//TEST_CASE("Write file and read back", "[datasource]") {
-//    vx_session *session = vx_session_new();
-//    TempPath path = write_sample(session, fs::current_path() / "write-read.vortex");
-//    vx_error *error = nullptr;
-//
-//    vx_data_source_options ds_options = {};
-//    ds_options.files = path.c_str();
-//
-//    const vx_data_source *ds = vx_data_source_new(session, &ds_options, &error);
-//    require_no_error(error);
-//    REQUIRE(ds != nullptr);
-//
-//    vx_scan *scan = vx_data_source_scan(ds, nullptr, nullptr, &error);
-//    require_no_error(error);
-//    REQUIRE(scan != nullptr);
-//    vx_scan_free(scan);
-//
-//    vx_scan_options scan_options = {};
-//    scan = vx_data_source_scan(ds, &scan_options, nullptr, &error);
-//    require_no_error(error);
-//    REQUIRE(scan != nullptr);
-//    vx_scan_free(scan);
-//
-//    vx_data_source_free(ds);
-//    vx_session_free(session);
-//}
+void verify_sample_array(const vx_array *array) {
+    REQUIRE(vx_array_len(array) == SAMPLE_ROWS);
+    REQUIRE(vx_array_has_dtype(array, DTYPE_STRUCT));
+
+    const vx_struct_fields *fields = vx_dtype_struct_dtype(vx_array_dtype(array));
+    size_t len = vx_struct_fields_nfields(fields);
+    REQUIRE(len == 2);
+
+    const vx_dtype *age_dtype = vx_struct_fields_field_dtype(fields, 0);
+    REQUIRE(vx_dtype_get_variant(age_dtype) == DTYPE_PRIMITIVE);
+    REQUIRE(vx_dtype_primitive_ptype(age_dtype) == PTYPE_U8);
+    vx_dtype_free(age_dtype);
+    const vx_string *age_name = vx_struct_fields_field_name(fields, 0);
+    REQUIRE(to_string_view(age_name) == "age");
+
+    const vx_dtype *height_dtype = vx_struct_fields_field_dtype(fields, 1);
+    REQUIRE(vx_dtype_get_variant(height_dtype) == DTYPE_PRIMITIVE);
+    REQUIRE(vx_dtype_primitive_ptype(height_dtype) == PTYPE_U16);
+    vx_dtype_free(height_dtype);
+    const vx_string *height_name = vx_struct_fields_field_name(fields, 1);
+    REQUIRE(to_string_view(height_name) == "height");
+
+    vx_error *error = nullptr;
+    vx_validity validity = {};
+    vx_array_get_validity(array, &validity, &error);
+    require_no_error(error);
+    REQUIRE(validity.type == VX_VALIDITY_NON_NULLABLE);
+
+    const vx_array *age_field = vx_array_get_field(array, 0, &error);
+    require_no_error(error);
+    REQUIRE(vx_array_has_dtype(age_field, DTYPE_PRIMITIVE));
+    REQUIRE(vx_dtype_primitive_ptype(vx_array_dtype(age_field)) == PTYPE_U8);
+    REQUIRE(vx_array_len(age_field) == SAMPLE_ROWS);
+    for (size_t i = 0; i < SAMPLE_ROWS; ++i) {
+        REQUIRE(vx_array_get_u8(age_field, i) == i);
+    }
+    vx_array_free(age_field);
+
+    const vx_array *height_field = vx_array_get_field(array, 1, &error);
+    require_no_error(error);
+    REQUIRE(vx_array_has_dtype(height_field, DTYPE_PRIMITIVE));
+    REQUIRE(vx_dtype_primitive_ptype(vx_array_dtype(height_field)) == PTYPE_U16);
+    REQUIRE(vx_array_len(height_field) == SAMPLE_ROWS);
+    for (size_t i = 0; i < SAMPLE_ROWS; ++i) {
+        REQUIRE(vx_array_get_u16(height_field, i) > 0);
+    }
+    vx_array_free(height_field);
+
+    REQUIRE(vx_array_get_field(array, 2, &error) == nullptr);
+    REQUIRE(error != nullptr);
+    vx_error_free(error);
+}
+
+TEST_CASE("Requesting scans", "[datasource]") {
+    vx_session *session = vx_session_new();
+    TempPath path = write_sample(session, fs::current_path() / "write-file2.vortex");
+
+    vx_data_source_options ds_options = {};
+    ds_options.files = path.c_str();
+
+    vx_error *error = nullptr;
+    const vx_data_source *ds = vx_data_source_new(session, &ds_options, &error);
+    require_no_error(error);
+    REQUIRE(ds != nullptr);
+
+    vx_scan *scan = vx_data_source_scan(ds, nullptr, nullptr, &error);
+    require_no_error(error);
+    REQUIRE(scan != nullptr);
+    vx_scan_free(scan);
+
+    vx_scan_options options = {};
+    options.max_threads = 1;
+    scan = vx_data_source_scan(ds, &options, nullptr, &error);
+    require_no_error(error);
+    REQUIRE(scan != nullptr);
+    vx_scan_free(scan);
+
+    vx_data_source_free(ds);
+    vx_session_free(session);
+}
+
+TEST_CASE("Basic scan", "[datasource]") {
+    vx_session *session = vx_session_new();
+    TempPath path = write_sample(session, fs::current_path() / "basic-scan.vortex");
+    vx_error *error = nullptr;
+
+    vx_data_source_options ds_options = {};
+    ds_options.files = path.c_str();
+
+    const vx_data_source *ds = vx_data_source_new(session, &ds_options, &error);
+    require_no_error(error);
+    REQUIRE(ds != nullptr);
+
+    vx_scan *scan = vx_data_source_scan(ds, nullptr, nullptr, &error);
+    require_no_error(error);
+    REQUIRE(scan != nullptr);
+
+    vx_partition *partition = vx_scan_next(scan, &error);
+    require_no_error(error);
+
+    vx_estimate estimate = {};
+    vx_partition_row_count(partition, &estimate, &error);
+    require_no_error(error);
+    REQUIRE(estimate.type == VX_ESTIMATE_EXACT);
+    REQUIRE(estimate.estimate == SAMPLE_ROWS);
+
+    REQUIRE(vx_scan_next(scan, &error) == nullptr);
+    require_no_error(error);
+
+    const vx_array *array = vx_partition_next(partition, &error);
+    require_no_error(error);
+    REQUIRE(array != nullptr);
+
+    REQUIRE(vx_partition_next(partition, &error) == nullptr);
+    require_no_error(error);
+
+    verify_sample_array(array);
+
+    vx_array_free(array);
+    vx_partition_free(partition);
+    vx_scan_free(scan);
+
+    vx_data_source_free(ds);
+    vx_session_free(session);
+}