diff --git a/.github/workflows/build-and-test-common.yml b/.github/workflows/build-and-test-common.yml index 209b7b8aa..9cac5b684 100644 --- a/.github/workflows/build-and-test-common.yml +++ b/.github/workflows/build-and-test-common.yml @@ -305,6 +305,7 @@ jobs: python3 -m pytest --cov=neug --cov-report=term --exitfirst --cov-append --cov-config=.coveragerc -sv tests/test_query.py python3 -m pytest --cov=neug --cov-report=term --exitfirst --cov-append --cov-config=.coveragerc -sv tests/test_tp_service.py python3 -m pytest --cov=neug --cov-report=term --exitfirst --cov-append --cov-config=.coveragerc -sv tests/test_load.py + python3 -m pytest --cov=neug --cov-report=term --exitfirst --cov-append --cov-config=.coveragerc -sv tests/test_load_array.py python3 -m pytest --cov=neug --cov-report=term --exitfirst --cov-append --cov-config=.coveragerc -sv tests/test_export.py python3 -m pytest --cov=neug --cov-report=term --exitfirst --cov-append --cov-config=.coveragerc -sv tests/test_copy_temp.py diff --git a/.github/workflows/neug-extension-test.yml b/.github/workflows/neug-extension-test.yml index f7f49a3be..5c152795e 100644 --- a/.github/workflows/neug-extension-test.yml +++ b/.github/workflows/neug-extension-test.yml @@ -140,6 +140,7 @@ jobs: export FLEX_DATA_DIR=${GITHUB_WORKSPACE}/example_dataset/comprehensive_graph GLOG_v=10 ${GITHUB_WORKSPACE}/build/tools/utils/bulk_loader -g ../../example_dataset/comprehensive_graph/graph.yaml -l ../../example_dataset/comprehensive_graph/import.yaml -d /tmp/comprehensive_graph NEUG_RUN_EXTENSION_TESTS=true python3 -m pytest -sv tests/test_load.py -k "parquet or httpfs" + NEUG_RUN_EXTENSION_TESTS=true python3 -m pytest -sv tests/test_load_array.py -k "parquet" NEUG_RUN_EXTENSION_TESTS=true python3 -m pytest -sv tests/test_export.py -k "parquet or httpfs" NEUG_RUN_EXTENSION_TESTS=true python3 -m pytest -sv tests/test_db_import_export.py -k "copy_from_parquet" NEUG_RUN_EXTENSION_TESTS=true python3 -m pytest -sv tests/test_sniffer.py -k "parquet" diff --git a/doc/source/data_io/import_data.md b/doc/source/data_io/import_data.md index abe9c042e..a27da5218 100644 --- a/doc/source/data_io/import_data.md +++ b/doc/source/data_io/import_data.md @@ -240,7 +240,25 @@ The following options control how CSV files are parsed during `COPY FROM`. These | `quoting` | bool | `true` | Whether to enable quote processing | | `escaping` | bool | `true` | Whether to enable escape character processing | -> **Array limitation:** `COPY FROM` CSV currently does not auto-detect or directly materialize fixed-size `ARRAY` columns from bracketed CSV fields such as `"[1,2,3]"`. Use Cypher literals/parameters for array properties, or a typed non-CSV ingestion path when available. +`COPY FROM` supports loading `ARRAY` data from CSV, +JSON, and Parquet files. NeuG does not implicitly convert input values to +`ARRAY`; use `LOAD FROM` and explicitly cast the column to a fixed-size +`ARRAY` type. + +For example, given a `person_array.csv` file with an array column: + +```csv +id,name,address +1,Alice,"[Beijing,Hangzhou,Shanghai]" +2,Bob,"[London,Paris,Berlin]" +``` + +```cypher +COPY Person FROM ( + LOAD FROM "person_array.csv" (delim=',') + RETURN id, name, CAST(address, 'STRING[3]') AS addresses +); +``` ### JSON/JSONL diff --git a/doc/source/data_io/load_data.md b/doc/source/data_io/load_data.md index db888b4c4..91c870203 100644 --- a/doc/source/data_io/load_data.md +++ b/doc/source/data_io/load_data.md @@ -46,7 +46,23 @@ LOAD FROM "person.csv" (delim=',', header=true) RETURN name, age; ``` -> **Array limitation:** `LOAD FROM` CSV currently does not parse CSV fields into fixed-size `ARRAY` values. A quoted field such as `"[1,2,3]"` is read as a `STRING`, and casting CSV columns to `INT64[3]`/other fixed-size array targets inside `LOAD FROM` is not supported yet. +NeuG supports reading `ARRAY` data from CSV, JSON, +and Parquet files. NeuG does not implicitly convert input values to `ARRAY`; +explicitly cast the column to a fixed-size `ARRAY` type in the `RETURN` +clause. + +For example, given a `person_array.csv` file with an array column: + +```csv +id,name,address +1,Alice,"[Beijing,Hangzhou,Shanghai]" +2,Bob,"[London,Paris,Berlin]" +``` + +```cypher +LOAD FROM "person_array.csv" (delim=',') +RETURN id, name, CAST(address, 'STRING[3]') AS addresses; +``` ### JSON / JSONL diff --git a/extension/parquet/src/arrow_column.cc b/extension/parquet/src/arrow_column.cc index 873e5adbb..4d5559cd2 100644 --- a/extension/parquet/src/arrow_column.cc +++ b/extension/parquet/src/arrow_column.cc @@ -16,13 +16,181 @@ #include "parquet/arrow_column.h" #include +#include #include +#include "neug/common/columns/columns_utils.h" #include "neug/common/columns/value_columns.h" +#include "neug/common/types/value.h" #include "neug/utils/exception/exception.h" namespace neug { +static int64_t arrow_time_to_milliseconds(int64_t value, + arrow::TimeUnit::type unit) { + switch (unit) { + case arrow::TimeUnit::SECOND: + return value * Interval::MSECS_PER_SEC; + case arrow::TimeUnit::MILLI: + return value; + case arrow::TimeUnit::MICRO: + return value / Interval::MICROS_PER_MSEC; + case arrow::TimeUnit::NANO: + return value / (Interval::MICROS_PER_MSEC * 1000); + } + THROW_NOT_SUPPORTED_EXCEPTION("Unsupported Arrow time unit"); +} + +// Keep this mapping in sync with arrow_value_at when adding Arrow types. +static DataType arrow_type_to_neug_type(const arrow::DataType& type) { + switch (type.id()) { + case arrow::Type::BOOL: + return DataType::BOOLEAN; + case arrow::Type::INT32: + return DataType::INT32; + case arrow::Type::INT64: + return DataType::INT64; + case arrow::Type::UINT32: + return DataType::UINT32; + case arrow::Type::UINT64: + return DataType::UINT64; + case arrow::Type::FLOAT: + return DataType::FLOAT; + case arrow::Type::DOUBLE: + return DataType::DOUBLE; + case arrow::Type::STRING: + case arrow::Type::LARGE_STRING: + return DataType::VARCHAR; + case arrow::Type::DATE32: + case arrow::Type::DATE64: + return DataType::DATE; + case arrow::Type::TIMESTAMP: + return DataType::TIMESTAMP_MS; + case arrow::Type::DURATION: + return DataType::INTERVAL; + case arrow::Type::FIXED_SIZE_LIST: { + const auto& array_type = static_cast(type); + return DataType::Array(arrow_type_to_neug_type(*array_type.value_type()), + array_type.list_size()); + } + case arrow::Type::LIST: + case arrow::Type::LARGE_LIST: + THROW_NOT_SUPPORTED_EXCEPTION( + "Parquet LIST is not supported as ARRAY. Specify a fixed-size Arrow " + "list / NeuG ARRAY type instead."); + default: + THROW_NOT_SUPPORTED_EXCEPTION("Unsupported arrow type: " + type.ToString()); + } +} + +static Value arrow_value_at(const arrow::Array& array, int64_t index, + const DataType& type) { + if (array.IsNull(index)) { + return Value(type); + } + switch (array.type_id()) { + case arrow::Type::BOOL: + return Value::BOOLEAN( + static_cast(array).Value(index)); + case arrow::Type::INT32: + return Value::INT32( + static_cast(array).Value(index)); + case arrow::Type::INT64: + return Value::INT64( + static_cast(array).Value(index)); + case arrow::Type::UINT32: + return Value::UINT32( + static_cast(array).Value(index)); + case arrow::Type::UINT64: + return Value::UINT64( + static_cast(array).Value(index)); + case arrow::Type::FLOAT: + return Value::FLOAT( + static_cast(array).Value(index)); + case arrow::Type::DOUBLE: + return Value::DOUBLE( + static_cast(array).Value(index)); + case arrow::Type::STRING: + return Value::STRING(std::string( + static_cast(array).GetView(index))); + case arrow::Type::LARGE_STRING: + return Value::STRING(std::string( + static_cast(array).GetView(index))); + case arrow::Type::DATE32: { + Date date; + date.from_num_days( + static_cast(array).Value(index)); + return Value::DATE(date); + } + case arrow::Type::DATE64: { + constexpr int64_t MILLIS_PER_DAY = 24LL * 60 * 60 * 1000; + Date date; + date.from_num_days(static_cast( + static_cast(array).Value(index) / + MILLIS_PER_DAY)); + return Value::DATE(date); + } + case arrow::Type::TIMESTAMP: { + const auto value = + static_cast(array).Value(index); + const auto& timestamp_type = + static_cast(*array.type()); + return Value::TIMESTAMPMS( + DateTime(arrow_time_to_milliseconds(value, timestamp_type.unit()))); + } + case arrow::Type::DURATION: { + const auto value = + static_cast(array).Value(index); + const auto& duration_type = + static_cast(*array.type()); + Interval interval; + interval.from_mill_seconds( + arrow_time_to_milliseconds(value, duration_type.unit())); + return Value::INTERVAL(interval); + } + case arrow::Type::FIXED_SIZE_LIST: { + const auto& list = static_cast(array); + const auto& array_type = + static_cast(*array.type()); + const auto child_type = ArrayType::GetChildType(type); + std::vector values; + values.reserve(array_type.list_size()); + const auto offset = list.value_offset(index); + for (int32_t i = 0; i < array_type.list_size(); ++i) { + values.push_back(arrow_value_at(*list.values(), offset + i, child_type)); + } + return Value::ARRAY(type, std::move(values)); + } + default: + THROW_NOT_SUPPORTED_EXCEPTION("Unsupported arrow type: " + + array.type()->ToString()); + } +} + +static std::shared_ptr convert_fixed_size_list_arrays( + const std::vector>& arrays) { + const auto type = arrow_type_to_neug_type(*arrays.front()->type()); + auto builder = ColumnsUtils::create_builder(type); + size_t size = 0; + for (const auto& array : arrays) { + size += array->length(); + } + builder->reserve(size); + for (const auto& array : arrays) { + if (!array->type()->Equals(*arrays.front()->type())) { + THROW_SCHEMA_MISMATCH("Parquet ARRAY chunks have different types"); + } + for (int64_t i = 0; i < array->length(); ++i) { + if (array->IsNull(i)) { + builder->push_back_null(); + continue; + } + builder->push_back_elem(arrow_value_at(*array, i, type)); + } + } + return builder->finish(); +} + /// Convert numeric arrow arrays directly to ValueColumn. template static std::shared_ptr convert_numeric_arrays( @@ -89,7 +257,11 @@ static std::shared_ptr convert_date64_arrays( if (typed->IsNull(j)) { builder.push_back_null(); } else { - builder.push_back_opt(Date(typed->Value(j))); + constexpr int64_t MILLIS_PER_DAY = 24LL * 60 * 60 * 1000; + Date date; + date.from_num_days( + static_cast(typed->Value(j) / MILLIS_PER_DAY)); + builder.push_back_opt(date); } } } @@ -106,7 +278,10 @@ static std::shared_ptr convert_timestamp_arrays( if (typed->IsNull(j)) { builder.push_back_null(); } else { - builder.push_back_opt(DateTime(typed->Value(j))); + const auto& timestamp_type = + static_cast(*typed->type()); + builder.push_back_opt(DateTime(arrow_time_to_milliseconds( + typed->Value(j), timestamp_type.unit()))); } } } @@ -144,6 +319,13 @@ std::shared_ptr arrow_arrays_to_value_column( return convert_date64_arrays(arrays); case arrow::Type::TIMESTAMP: return convert_timestamp_arrays(arrays); + case arrow::Type::FIXED_SIZE_LIST: + return convert_fixed_size_list_arrays(arrays); + case arrow::Type::LIST: + case arrow::Type::LARGE_LIST: + THROW_NOT_SUPPORTED_EXCEPTION( + "Parquet LIST columns are not supported; use a fixed-size ARRAY " + "type such as FLOAT[3]"); default: THROW_NOT_SUPPORTED_EXCEPTION("Unsupported arrow type: " + arrow_type->ToString()); diff --git a/include/neug/storages/loader/loader_utils.h b/include/neug/storages/loader/loader_utils.h index 32b33aab9..5b0eb5c1f 100644 --- a/include/neug/storages/loader/loader_utils.h +++ b/include/neug/storages/loader/loader_utils.h @@ -99,8 +99,8 @@ void fillVertexReaderMeta(label_t v_label, const std::string& v_label_name, const std::string& v_file, const LoadingConfig& loading_config, const std::vector& vertex_property_names, - const std::vector& vertex_property_types, - DataTypeId pk_type, const std::string& pk_name, + const std::vector& vertex_property_types, + DataType pk_type, const std::string& pk_name, size_t pk_ind, CsvReadConfig& config); void fillEdgeReaderMeta(label_t src_label_id, label_t dst_label_id, @@ -108,8 +108,8 @@ void fillEdgeReaderMeta(label_t src_label_id, label_t dst_label_id, const std::string& e_file, const LoadingConfig& loading_config, const std::vector& edge_property_names, - const std::vector& edge_property_types, - DataTypeId src_pk_type, DataTypeId dst_pk_type, + const std::vector& edge_property_types, + DataType src_pk_type, DataType dst_pk_type, CsvReadConfig& config); void set_properties_from_context_column( diff --git a/src/storages/loader/csv_property_graph_loader.cc b/src/storages/loader/csv_property_graph_loader.cc index b61a81071..ff862de3e 100644 --- a/src/storages/loader/csv_property_graph_loader.cc +++ b/src/storages/loader/csv_property_graph_loader.cc @@ -32,12 +32,12 @@ CSVPropertyGraphLoader::createVertexChunkSupplier( DataType pk_type, const std::string& pk_name, int pk_ind, const LoadingConfig& loading_config, int thread_id) const { auto vertex_property_names = schema_.get_vertex_property_names(v_label); - auto vertex_property_types = schema_.get_vertex_properties_id(v_label); + auto vertex_property_types = schema_.get_vertex_properties(v_label); CsvReadConfig config; fillVertexReaderMeta(v_label, v_label_name, v_file, loading_config, - vertex_property_names, vertex_property_types, - pk_type.id(), pk_name, pk_ind, config); + vertex_property_names, vertex_property_types, pk_type, + pk_name, pk_ind, config); return std::make_shared(v_file, std::move(config)); } @@ -49,7 +49,7 @@ CSVPropertyGraphLoader::createEdgeChunkSupplier( auto edge_property_names = schema_.get_edge_property_names(src_label_id, dst_label_id, e_label_id); auto edge_property_types = - schema_.get_edge_properties_id(src_label_id, dst_label_id, e_label_id); + schema_.get_edge_properties(src_label_id, dst_label_id, e_label_id); auto src_pk_type = std::get<0>(schema_.get_vertex_primary_key(src_label_id)[0]); auto dst_pk_type = @@ -58,7 +58,7 @@ CSVPropertyGraphLoader::createEdgeChunkSupplier( fillEdgeReaderMeta(src_label_id, dst_label_id, e_label_id, schema_.get_edge_label_name(e_label_id), e_file, loading_config_, edge_property_names, edge_property_types, - src_pk_type.id(), dst_pk_type.id(), config); + src_pk_type, dst_pk_type, config); return std::make_shared(e_file, std::move(config)); } diff --git a/src/storages/loader/loader_utils.cc b/src/storages/loader/loader_utils.cc index 4c2d0e8ba..6636f6105 100644 --- a/src/storages/loader/loader_utils.cc +++ b/src/storages/loader/loader_utils.cc @@ -251,6 +251,156 @@ T parse_direct(std::string_view token, } } +std::string_view trim_array_token(std::string_view token) { + size_t begin = 0; + while (begin < token.size() && + std::isspace(static_cast(token[begin]))) { + ++begin; + } + size_t end = token.size(); + while (end > begin && + std::isspace(static_cast(token[end - 1]))) { + --end; + } + return token.substr(begin, end - begin); +} + +std::string unquote_array_string_token(std::string_view token) { + auto trimmed = trim_array_token(token); + if (trimmed.size() < 2) { + return std::string(trimmed); + } + auto quote = trimmed.front(); + if ((quote != '\'' && quote != '"') || trimmed.back() != quote) { + return std::string(trimmed); + } + + std::string unquoted; + unquoted.reserve(trimmed.size() - 2); + for (size_t i = 1; i + 1 < trimmed.size(); ++i) { + if (trimmed[i] == '\\' && i + 1 < trimmed.size() - 1) { + unquoted.push_back(trimmed[++i]); + } else { + unquoted.push_back(trimmed[i]); + } + } + return unquoted; +} + +std::vector split_array_elements(std::string_view input) { + std::vector elements; + auto trimmed = trim_array_token(input); + if (trimmed.empty()) { + return elements; + } + + int depth = 0; + char quote = '\0'; + bool escaping = false; + size_t element_start = 0; + for (size_t i = 0; i < trimmed.size(); ++i) { + auto c = trimmed[i]; + if (quote != '\0') { + if (escaping) { + escaping = false; + } else if (c == '\\') { + escaping = true; + } else if (c == quote) { + quote = '\0'; + } + continue; + } + if (c == '\'' || c == '"') { + quote = c; + } else if (c == '[') { + ++depth; + } else if (c == ']') { + if (depth == 0) { + THROW_CONVERSION_EXCEPTION("Unexpected ']' in array value: " + + std::string(input)); + } + --depth; + } else if (c == ',' && depth == 0) { + elements.push_back( + trim_array_token(trimmed.substr(element_start, i - element_start))); + element_start = i + 1; + } + } + if (quote != '\0') { + THROW_CONVERSION_EXCEPTION("Unterminated quoted string in array value: " + + std::string(input)); + } + if (depth != 0) { + THROW_CONVERSION_EXCEPTION("Unbalanced brackets in array value: " + + std::string(input)); + } + elements.push_back(trim_array_token(trimmed.substr(element_start))); + return elements; +} + +Value parse_array_element(std::string_view token, const DataType& data_type, + const std::unordered_set& true_values, + const std::unordered_set& false_values) { + switch (data_type.id()) { + case DataTypeId::kInt32: + return Value::INT32( + parse_direct(token, true_values, false_values)); + case DataTypeId::kInt64: + return Value::INT64( + parse_direct(token, true_values, false_values)); + case DataTypeId::kUInt32: + return Value::UINT32( + parse_direct(token, true_values, false_values)); + case DataTypeId::kUInt64: + return Value::UINT64( + parse_direct(token, true_values, false_values)); + case DataTypeId::kFloat: + return Value::FLOAT(parse_direct(token, true_values, false_values)); + case DataTypeId::kDouble: + return Value::DOUBLE( + parse_direct(token, true_values, false_values)); + case DataTypeId::kBoolean: + return Value::BOOLEAN(parse_direct(token, true_values, false_values)); + case DataTypeId::kDate: + return Value::DATE(parse_direct(token, true_values, false_values)); + case DataTypeId::kTimestampMs: + return Value::TIMESTAMPMS( + parse_direct(token, true_values, false_values)); + case DataTypeId::kInterval: + return Value::INTERVAL( + parse_direct(token, true_values, false_values)); + case DataTypeId::kVarchar: + return Value::STRING(unquote_array_string_token(token)); + case DataTypeId::kArray: { + auto trimmed = trim_array_token(token); + if (trimmed.size() < 2 || trimmed.front() != '[' || trimmed.back() != ']') { + THROW_CONVERSION_EXCEPTION("Expected array value for type " + + data_type.ToString() + ": " + + std::string(trimmed)); + } + const auto& child_type = ArrayType::GetChildType(data_type); + const auto expected_size = ArrayType::GetNumElements(data_type); + auto elements = split_array_elements(trimmed.substr(1, trimmed.size() - 2)); + if (elements.size() != expected_size) { + THROW_CONVERSION_EXCEPTION("ARRAY value length mismatch for type " + + data_type.ToString() + ": expected " + + std::to_string(expected_size) + ", got " + + std::to_string(elements.size())); + } + std::vector values; + values.reserve(elements.size()); + for (const auto& element : elements) { + values.push_back( + parse_array_element(element, child_type, true_values, false_values)); + } + return Value::ARRAY(data_type, std::move(values)); + } + default: + THROW_NOT_SUPPORTED_EXCEPTION("Unsupported ARRAY element type: " + + data_type.ToString()); + } +} + /// Shared parsing context (constant across all columns in a chunk). struct CsvParseCtx { const std::unordered_set& null_values; @@ -317,10 +467,42 @@ void append_typed_impl(const FieldAppender& app, csv::CSVField field, } } +void append_array_impl(const FieldAppender& app, csv::CSVField field, + const CsvParseCtx& ctx, int64_t row_number) { + auto* builder = static_cast(app.builder); + if (field.is_null()) { + builder->push_back_null(); + return; + } + std::string_view token = static_cast(field); + if (contains_sv(ctx.null_values, token)) { + builder->push_back_null(); + return; + } + std::string unescaped; + std::string_view effective_token = token; + if (ctx.escaping) { + unescaped = unescape_token_sv(token, ctx.escape_char); + effective_token = unescaped; + } + try { + builder->push_back_elem(parse_array_element( + effective_token, *app.type, ctx.true_values, ctx.false_values)); + } catch (const std::exception& error) { + THROW_CONVERSION_EXCEPTION( + "Failed to parse CSV field, file=" + ctx.file_path + + ", row=" + std::to_string(row_number) + ", column=" + *app.column_name + + ", type=" + app.type->ToString() + ", value='" + + std::string(effective_token) + "', reason=" + error.what()); + } +} + /// Create a FieldAppender for the given data type. FieldAppender make_appender(const DataType& type, void* builder, const std::string* column_name) { switch (type.id()) { + case DataTypeId::kArray: + return {builder, column_name, &type, &append_array_impl}; #define MAKE_APPENDER(enum_val, cpp_type) \ case DataTypeId::enum_val: \ return {builder, column_name, &type, &append_typed_impl}; @@ -1057,9 +1239,8 @@ void fillVertexReaderMeta( label_t v_label, const std::string& v_label_name, const std::string& v_file, const LoadingConfig& loading_config, const std::vector& vertex_property_names, - const std::vector& vertex_edge_property_types, - DataTypeId pk_type, const std::string& pk_name, size_t pk_ind, - CsvReadConfig& config) { + const std::vector& vertex_edge_property_types, DataType pk_type, + const std::string& pk_name, size_t pk_ind, CsvReadConfig& config) { CHECK(vertex_edge_property_types.size() == vertex_property_names.size()); put_boolean_option(config); @@ -1140,9 +1321,9 @@ void fillVertexReaderMeta( } VLOG(10) << "vertex_label: " << v_label_name << " property_name: " << property_name - << " property_type: " << property_type << " ind: " << ind; - config.column_types.insert( - {included_col_names[ind], DataType(property_type)}); + << " property_type: " << property_type.ToString() + << " ind: " << ind; + config.column_types.insert({included_col_names[ind], property_type}); } { size_t ind = mapped_property_names.size(); @@ -1158,7 +1339,7 @@ void fillVertexReaderMeta( " does not exist in the vertex column mapping, please " "check your configuration"); } - config.column_types.insert({included_col_names[ind], DataType(pk_type)}); + config.column_types.insert({included_col_names[ind], pk_type}); } } @@ -1167,8 +1348,8 @@ void fillEdgeReaderMeta(label_t src_label_id, label_t dst_label_id, const std::string& e_file, const LoadingConfig& loading_config, const std::vector& edge_property_names, - const std::vector& edge_property_types, - DataTypeId src_pk_type, DataTypeId dst_pk_type, + const std::vector& edge_property_types, + DataType src_pk_type, DataType dst_pk_type, CsvReadConfig& config) { CHECK(edge_property_types.size() == edge_property_names.size()); put_boolean_option(config); @@ -1265,9 +1446,9 @@ void fillEdgeReaderMeta(label_t src_label_id, label_t dst_label_id, } VLOG(10) << "edge_label: " << edge_label_name << " property_name: " << property_name - << " property_type: " << property_type << " ind: " << ind; - config.column_types.insert( - {included_col_names[ind + 2], DataType(property_type)}); + << " property_type: " << property_type.ToString() + << " ind: " << ind; + config.column_types.insert({included_col_names[ind + 2], property_type}); } { auto src_dst_cols = @@ -1279,10 +1460,8 @@ void fillEdgeReaderMeta(label_t src_label_id, label_t dst_label_id, src_col_ind < static_cast(config.column_names.size())); CHECK(dst_col_ind >= 0 && dst_col_ind < static_cast(config.column_names.size())); - config.column_types.insert( - {config.column_names[src_col_ind], DataType(src_pk_type)}); - config.column_types.insert( - {config.column_names[dst_col_ind], DataType(dst_pk_type)}); + config.column_types.insert({config.column_names[src_col_ind], src_pk_type}); + config.column_types.insert({config.column_names[dst_col_ind], dst_pk_type}); VLOG(10) << "Column types: "; for (const auto& iter : config.column_types) { @@ -1367,6 +1546,18 @@ void set_properties_from_context_column( vids); break; } + case DataTypeId::kArray: { + for (size_t k = 0; k < vids.size(); ++k) { + if (vids[k] >= std::numeric_limits::max()) { + continue; + } + auto value = ctx_col->get_elem(k); + if (!value.IsNull()) { + col->set_any(vids[k], value, true); + } + } + break; + } default: THROW_NOT_SUPPORTED_EXCEPTION( "set_properties_from_context_column: unsupported type " + diff --git a/src/utils/io/read/json/json_reader.cc b/src/utils/io/read/json/json_reader.cc index 7bbd1cd8c..fa7d0f933 100644 --- a/src/utils/io/read/json/json_reader.cc +++ b/src/utils/io/read/json/json_reader.cc @@ -148,6 +148,26 @@ Value parse_json_value(const rapidjson::Value& value, return Value::STRING(value.GetString()); } return Value::STRING(rapidjson_stringify(value)); + case DataTypeId::kArray: { + if (!value.IsArray()) { + THROW_CONVERSION_EXCEPTION("Expected JSON array for ARRAY type: " + + data_type.ToString()); + } + const auto expected_size = ArrayType::GetNumElements(data_type); + if (value.Size() != expected_size) { + THROW_CONVERSION_EXCEPTION("ARRAY value length mismatch for type " + + data_type.ToString() + ": expected " + + std::to_string(expected_size) + ", got " + + std::to_string(value.Size())); + } + const auto& child_type = ArrayType::GetChildType(data_type); + std::vector values; + values.reserve(value.Size()); + for (const auto& item : value.GetArray()) { + values.push_back(parse_json_value(item, child_type)); + } + return Value::ARRAY(data_type, std::move(values)); + } default: if (value.IsString()) { return Value::STRING(value.GetString()); diff --git a/tests/utils/json_test.cc b/tests/utils/json_test.cc index 46b8fd395..d526bf189 100644 --- a/tests/utils/json_test.cc +++ b/tests/utils/json_test.cc @@ -75,6 +75,16 @@ class JsonTest : public ::testing::Test { return type; } + std::shared_ptr<::common::DataType> createInt64ArrayType( + uint32_t fixed_length) { + auto type = std::make_shared<::common::DataType>(); + auto* array = type->mutable_array(); + array->set_fixed_length(fixed_length); + array->mutable_component_type()->set_primitive_type( + ::common::PrimitiveType::DT_SIGNED_INT64); + return type; + } + std::shared_ptr createSharedState( const std::string& jsonFile, const std::vector& columnNames, const std::vector>& columnTypes, @@ -140,5 +150,72 @@ TEST_F(JsonTest, TestJsonArray) { EXPECT_DOUBLE_EQ(col2->get_elem(1).GetValue(), 30.0); } +TEST_F(JsonTest, TestJsonArrayColumn) { + createJsonFile("test_json_array_column.json", + "[{\"id\": 1, \"readings\": [1, 2, 3]}, " + "{\"id\": 2, \"readings\": [4, 5, 6]}]"); + auto sharedState = createSharedState( + "test_json_array_column.json", {"id", "readings"}, + {createUInt32Type(), createInt64ArrayType(3)}, {{"batch_read", "false"}}); + auto reader = createJsonReader(sharedState); + auto localState = std::make_shared(); + execution::Context ctx; + + reader->read(localState, ctx); + + EXPECT_EQ(ctx.col_num(), 2); + EXPECT_EQ(ctx.row_num(), 2); + + auto readings_col = ctx.chunk(0).columns()[1]; + ASSERT_EQ(readings_col->column_type(), ContextColumnType::kValue); + auto first = readings_col->get_elem(0); + const auto& first_values = ArrayValue::GetChildren(first); + ASSERT_EQ(first_values.size(), 3); + EXPECT_EQ(first_values[0].GetValue(), 1); + EXPECT_EQ(first_values[1].GetValue(), 2); + EXPECT_EQ(first_values[2].GetValue(), 3); +} + +TEST_F(JsonTest, TestJsonArrayColumnLengthMismatch) { + createJsonFile("test_json_array_length_mismatch.json", + "[{\"id\": 1, \"readings\": [1, 2]}]"); + auto sharedState = createSharedState( + "test_json_array_length_mismatch.json", {"id", "readings"}, + {createUInt32Type(), createInt64ArrayType(3)}, {{"batch_read", "false"}}); + auto reader = createJsonReader(sharedState); + auto localState = std::make_shared(); + execution::Context ctx; + + try { + reader->read(localState, ctx); + FAIL() << "Expected an ARRAY length mismatch"; + } catch (const std::exception& error) { + EXPECT_NE( + std::string(error.what()) + .find("ARRAY value length mismatch for type INT64[3]: expected 3, " + "got 2"), + std::string::npos); + } +} + +TEST_F(JsonTest, TestJsonArrayColumnRejectsNonArray) { + createJsonFile("test_json_non_array.json", "[{\"id\": 1, \"readings\": 42}]"); + auto sharedState = createSharedState( + "test_json_non_array.json", {"id", "readings"}, + {createUInt32Type(), createInt64ArrayType(3)}, {{"batch_read", "false"}}); + auto reader = createJsonReader(sharedState); + auto localState = std::make_shared(); + execution::Context ctx; + + try { + reader->read(localState, ctx); + FAIL() << "Expected a non-array conversion error"; + } catch (const std::exception& error) { + EXPECT_NE(std::string(error.what()) + .find("Expected JSON array for ARRAY type: INT64[3]"), + std::string::npos); + } +} + } // namespace test } // namespace neug diff --git a/tools/python_bind/tests/test_data_io_docs.py b/tools/python_bind/tests/test_data_io_docs.py index 2fa5bf7f9..09a4f449a 100644 --- a/tools/python_bind/tests/test_data_io_docs.py +++ b/tools/python_bind/tests/test_data_io_docs.py @@ -210,8 +210,8 @@ def test_load_from_performance_options(self): records = list(result) assert len(records) == 4 - def test_load_from_csv_array_limitation(self): - """load_data.md: CSV fields are not parsed as fixed-size ARRAY values.""" + def test_load_from_csv_array_with_explicit_cast(self): + """load_data.md: CSV fields can be cast to fixed-size ARRAY values.""" csv_path = self.tmp_path / "sensor_arrays.csv" csv_path.write_text('id,readings\n1,"[1,2,3]"\n') @@ -223,12 +223,13 @@ def test_load_from_csv_array_limitation(self): ) assert rows == [[1, "[1,2,3]"]] - with pytest.raises(RuntimeError) as exc_info: + rows = list( self.conn.execute( f'LOAD FROM "{csv_path}" (header=true, delimiter=",") ' - "RETURN id, CAST(readings, 'INT64[3]');" + "RETURN id, CAST(readings, 'INT64[3]') AS readings;" ) - assert "not supported" in str(exc_info.value).lower() + ) + assert rows == [[1, [1, 2, 3]]] # ============================================================ @@ -429,19 +430,24 @@ def test_copy_from_parallel(self): res = self.conn.execute("MATCH (p:Person) RETURN count(p);") assert list(res)[0][0] == 4 - def test_copy_from_csv_array_limitation(self): - """import_data.md: COPY FROM CSV cannot materialize fixed-size ARRAY columns.""" + def test_copy_from_csv_array_with_explicit_cast(self): + """import_data.md: COPY an explicitly cast CSV ARRAY column.""" csv_path = self.tmp_path / "sensor_arrays.csv" csv_path.write_text('id,readings\n1,"[1,2,3]"\n') self.conn.execute( "CREATE NODE TABLE Sensor(" "id INT64, readings INT64[3], PRIMARY KEY(id));" ) - with pytest.raises(RuntimeError) as exc_info: - self.conn.execute( - f'COPY Sensor FROM "{csv_path}" (header=true, delimiter=",");' - ) - assert "Unsupported data type in CSV parser" in str(exc_info.value) + self.conn.execute( + f""" + COPY Sensor FROM ( + LOAD FROM "{csv_path}" (header=true, delimiter=",") + RETURN id, CAST(readings, 'INT64[3]') AS readings + ); + """ + ) + rows = list(self.conn.execute("MATCH (s:Sensor) RETURN s.id, s.readings;")) + assert rows == [[1, [1, 2, 3]]] def test_copy_from_batch_read(self): """import_data.md: COPY FROM with batch_read and batch_size options. diff --git a/tools/python_bind/tests/test_load_array.py b/tools/python_bind/tests/test_load_array.py new file mode 100644 index 000000000..9142f6b04 --- /dev/null +++ b/tools/python_bind/tests/test_load_array.py @@ -0,0 +1,641 @@ +#!/usr/bin/env python3 +# -*- coding: utf-8 -*- +# +# Copyright 2020 Alibaba Group Holding Limited. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# + +"""End-to-end LOAD FROM tests for fixed-size array columns.""" + +import json +import os +import shutil +import sys +from datetime import date +from datetime import datetime +from datetime import timedelta + +import pytest + +sys.path.append(os.path.join(os.path.dirname(__file__), "../")) + +from neug import Database + +EXTENSION_TESTS_ENABLED = os.environ.get("NEUG_RUN_EXTENSION_TESTS", "").lower() in ( + "1", + "true", + "yes", + "on", +) +extension_test = pytest.mark.skipif( + not EXTENSION_TESTS_ENABLED, + reason="Extension tests disabled by default; set NEUG_RUN_EXTENSION_TESTS=1 to enable.", +) + + +class TestLoadArray: + """Test cases for LOAD FROM CSV with CAST to array types.""" + + @pytest.fixture(autouse=True) + def setup(self, tmp_path): + """Setup test database and CSV directory.""" + self.db_dir = str(tmp_path / "test_load_array_db") + self.csv_dir = str(tmp_path / "csv_data") + self.json_dir = str(tmp_path / "json_data") + self.parquet_dir = str(tmp_path / "parquet_data") + shutil.rmtree(self.db_dir, ignore_errors=True) + os.makedirs(self.csv_dir, exist_ok=True) + os.makedirs(self.json_dir, exist_ok=True) + os.makedirs(self.parquet_dir, exist_ok=True) + self.db = Database(db_path=self.db_dir, mode="w") + self.conn = self.db.connect() + yield + self.conn.close() + self.db.close() + shutil.rmtree(self.db_dir, ignore_errors=True) + + def _write_csv(self, filename, content): + """Write a CSV file to the temp directory and return its path.""" + path = os.path.join(self.csv_dir, filename) + with open(path, "w", encoding="utf-8", newline="") as f: + f.write(content) + return path + + def _write_json(self, filename, data): + """Write a JSON file to the temp directory and return its path.""" + path = os.path.join(self.json_dir, filename) + with open(path, "w", encoding="utf-8") as f: + json.dump(data, f) + return path + + def _write_parquet(self, filename, columns): + """Write an Arrow table containing fixed-size list columns.""" + pa = pytest.importorskip("pyarrow") + pq = pytest.importorskip("pyarrow.parquet") + path = os.path.join(self.parquet_dir, filename) + pq.write_table(pa.table(columns), path) + return path + + def test_copy_csv_array_with_explicit_cast(self): + """COPY an ARRAY property through LOAD FROM with an explicit type.""" + csv_path = self._write_csv( + "person_array.csv", + "id,name,address\n" + '1,Alice,"[Beijing,Hangzhou,Shanghai]"\n' + '2,Bob,"[London,Paris,Berlin]"\n', + ) + self.conn.execute( + "CREATE NODE TABLE Person(" + "id INT64, name STRING, addresses STRING[3], PRIMARY KEY(id));" + ) + + self.conn.execute( + f""" + COPY Person FROM ( + LOAD FROM "{csv_path}" (delim=',') + RETURN id, name, CAST(address, 'STRING[3]') AS addresses + ) + """ + ) + + result = list( + self.conn.execute( + "MATCH (p:Person) " "RETURN p.id, p.name, p.addresses ORDER BY p.id" + ) + ) + assert result == [ + [1, "Alice", ["Beijing", "Hangzhou", "Shanghai"]], + [2, "Bob", ["London", "Paris", "Berlin"]], + ] + + @extension_test + def test_parquet_float_array(self): + """LOAD FROM Parquet with a FLOAT[3] fixed-size array.""" + pa = pytest.importorskip("pyarrow") + parquet_path = self._write_parquet( + "float_array.parquet", + { + "id": pa.array([1, 2], type=pa.int64()), + "values": pa.array( + [[1.5, 2.5, 3.5], [4.0, 5.0, 6.0]], + type=pa.list_(pa.float32(), 3), + ), + }, + ) + self.conn.execute("LOAD PARQUET") + result = list( + self.conn.execute( + f'LOAD FROM "{parquet_path}" ' + "RETURN id, CAST(values, 'FLOAT[3]') ORDER BY id" + ) + ) + + assert result == [ + [1, [1.5, 2.5, 3.5]], + [2, [4.0, 5.0, 6.0]], + ] + + @extension_test + def test_parquet_nested_string_array(self): + """LOAD FROM Parquet with a nested STRING[2][2] array.""" + pa = pytest.importorskip("pyarrow") + inner_type = pa.list_(pa.string(), 2) + parquet_path = self._write_parquet( + "nested_string_array.parquet", + { + "id": pa.array([1, 2], type=pa.int64()), + "values": pa.array( + [[["a", "b"], ["c", "d"]], [["w", "x"], ["y", "z"]]], + type=pa.list_(inner_type, 2), + ), + }, + ) + self.conn.execute("LOAD PARQUET") + result = list( + self.conn.execute( + f'LOAD FROM "{parquet_path}" ' + "RETURN id, CAST(values, 'STRING[2][2]') ORDER BY id" + ) + ) + + assert result == [ + [1, [["a", "b"], ["c", "d"]]], + [2, [["w", "x"], ["y", "z"]]], + ] + + @extension_test + def test_parquet_null_array(self): + """LOAD FROM Parquet preserves a null fixed-size array.""" + pa = pytest.importorskip("pyarrow") + parquet_path = self._write_parquet( + "null_array.parquet", + { + "id": pa.array([1, 2, 3], type=pa.int64()), + "values": pa.array( + [[1.0, 2.0, 3.0], None, [7.0, 8.0, 9.0]], + type=pa.list_(pa.float32(), 3), + ), + }, + ) + self.conn.execute("LOAD PARQUET") + result = list( + self.conn.execute( + f'LOAD FROM "{parquet_path}" ' + "RETURN id, CAST(values, 'FLOAT[3]') ORDER BY id" + ) + ) + + assert result == [ + [1, [1.0, 2.0, 3.0]], + [2, None], + [3, [7.0, 8.0, 9.0]], + ] + + @extension_test + def test_parquet_timestamp_units(self): + """Normalize scalar and ARRAY timestamps of every Arrow unit to ms.""" + pa = pytest.importorskip("pyarrow") + expected = datetime(2023, 6, 15, 12, 30, 45, 123000) + parquet_path = self._write_parquet( + "timestamp_units.parquet", + { + "seconds": pa.array( + [datetime(2023, 6, 15, 12, 30, 45)], pa.timestamp("s") + ), + "milliseconds": pa.array([expected], pa.timestamp("ms")), + "microseconds": pa.array([expected], pa.timestamp("us")), + "nanoseconds": pa.array([expected], pa.timestamp("ns")), + "times": pa.array( + [[expected, expected]], pa.list_(pa.timestamp("us"), 2) + ), + }, + ) + self.conn.execute("LOAD PARQUET") + + result = list( + self.conn.execute( + f'LOAD FROM "{parquet_path}" ' + "RETURN seconds, milliseconds, microseconds, nanoseconds, times" + ) + ) + + assert result == [ + [ + datetime(2023, 6, 15, 12, 30, 45), + expected, + expected, + expected, + [expected, expected], + ] + ] + + @extension_test + def test_parquet_array_types(self): + """Cover the scalar and nested ARRAY types exercised by CSV.""" + pa = pytest.importorskip("pyarrow") + parquet_path = self._write_parquet( + "array_types.parquet", + { + "id": pa.array([1], type=pa.int64()), + "floats": pa.array([[1.5, 2.5, 3.5]], pa.list_(pa.float32(), 3)), + "doubles": pa.array([[1.1, 2.2]], pa.list_(pa.float64(), 2)), + "int32s": pa.array([[10, 20, 30]], pa.list_(pa.int32(), 3)), + "int64s": pa.array( + [[100000000000, 200000000000]], pa.list_(pa.int64(), 2) + ), + "strings": pa.array( + [["hello", "world", "test"]], pa.list_(pa.string(), 3) + ), + "int_matrix": pa.array( + [[[1, 2], [3, 4]]], + pa.list_(pa.list_(pa.int32(), 2), 2), + ), + "string_matrix": pa.array( + [[["a", "b"], ["c", "d"]]], + pa.list_(pa.list_(pa.string(), 2), 2), + ), + "dates": pa.array( + [[date(1970, 1, 1), date(2023, 6, 15)]], + pa.list_(pa.date64(), 2), + ), + "times": pa.array( + [[datetime(1970, 1, 1), datetime(2023, 6, 15, 12, 30)]], + pa.list_(pa.timestamp("ms"), 2), + ), + "durations": pa.array( + [[timedelta(days=2), timedelta(hours=3)]], + pa.list_(pa.duration("ms"), 2), + ), + }, + ) + self.conn.execute("LOAD PARQUET") + result = list( + self.conn.execute( + f""" + LOAD FROM "{parquet_path}" + RETURN id, + CAST(floats, 'FLOAT[3]'), + CAST(doubles, 'DOUBLE[2]'), + CAST(int32s, 'INT32[3]'), + CAST(int64s, 'INT64[2]'), + CAST(strings, 'STRING[3]'), + CAST(int_matrix, 'INT32[2][2]'), + CAST(string_matrix, 'STRING[2][2]'), + CAST(dates, 'DATE[2]'), + CAST(times, 'TIMESTAMP[2]'), + CAST(durations, 'INTERVAL[2]') + """ + ) + ) + + assert result == [ + [ + 1, + [1.5, 2.5, 3.5], + [1.1, 2.2], + [10, 20, 30], + [100000000000, 200000000000], + ["hello", "world", "test"], + [[1, 2], [3, 4]], + [["a", "b"], ["c", "d"]], + [date(1970, 1, 1), date(2023, 6, 15)], + [datetime(1970, 1, 1), datetime(2023, 6, 15, 12, 30)], + ["2 days", "3 hours"], + ] + ] + + def test_cast_float_array(self): + """LOAD FROM CSV with CAST(col, 'FLOAT[3]').""" + csv_path = self._write_csv( + "float_array.csv", + "id|values\n" + "1|[1.5, 2.5, 3.5]\n" + "2|[4.0, 5.0, 6.0]\n" + "3|[0.0, 0.0, 0.0]\n", + ) + query = f""" + LOAD FROM "{csv_path}" (delim='|') + RETURN id, CAST(values, 'FLOAT[3]') + """ + result = list(self.conn.execute(query)) + assert len(result) == 3 + + # Row 1: [1.5, 2.5, 3.5] + assert len(result[0][1]) == 3 + assert abs(result[0][1][0] - 1.5) < 1e-5 + assert abs(result[0][1][1] - 2.5) < 1e-5 + assert abs(result[0][1][2] - 3.5) < 1e-5 + + # Row 2: [4.0, 5.0, 6.0] + assert len(result[1][1]) == 3 + assert abs(result[1][1][0] - 4.0) < 1e-5 + + # Row 3: [0.0, 0.0, 0.0] + assert result[2][1] == [0.0, 0.0, 0.0] + + def test_cast_double_array(self): + """LOAD FROM CSV with CAST(col, 'DOUBLE[2]').""" + csv_path = self._write_csv( + "double_array.csv", + "id|values\n" + "1|[1.111111111, 2.222222222]\n" + "2|[3.333333333, 4.444444444]\n", + ) + query = f""" + LOAD FROM "{csv_path}" (delim='|') + RETURN id, CAST(values, 'DOUBLE[2]') + """ + result = list(self.conn.execute(query)) + assert len(result) == 2 + + assert len(result[0][1]) == 2 + assert abs(result[0][1][0] - 1.111111111) < 1e-9 + assert abs(result[0][1][1] - 2.222222222) < 1e-9 + + assert len(result[1][1]) == 2 + assert abs(result[1][1][0] - 3.333333333) < 1e-9 + + def test_cast_int32_array(self): + """LOAD FROM CSV with CAST(col, 'INT32[3]').""" + csv_path = self._write_csv( + "int32_array.csv", + "id|nums\n" "1|[10, 20, 30]\n" "2|[-1, 0, 1]\n", + ) + query = f""" + LOAD FROM "{csv_path}" (delim='|') + RETURN id, CAST(nums, 'INT32[3]') + """ + result = list(self.conn.execute(query)) + assert len(result) == 2 + assert result[0][1] == [10, 20, 30] + assert result[1][1] == [-1, 0, 1] + + def test_cast_int64_array(self): + """LOAD FROM CSV with CAST(col, 'INT64[2]').""" + csv_path = self._write_csv( + "int64_array.csv", + "id|nums\n" "1|[100000000000, 200000000000]\n" "2|[0, 1]\n", + ) + query = f""" + LOAD FROM "{csv_path}" (delim='|') + RETURN id, CAST(nums, 'INT64[2]') + """ + result = list(self.conn.execute(query)) + assert len(result) == 2 + assert result[0][1] == [100000000000, 200000000000] + assert result[1][1] == [0, 1] + + def test_cast_string_array(self): + """LOAD FROM CSV with CAST(col, 'STRING[3]').""" + csv_path = self._write_csv( + "string_array.csv", + "id|tags\n" "1|[hello, world, test]\n" "2|[foo, bar, baz]\n", + ) + query = f""" + LOAD FROM "{csv_path}" (delim='|') + RETURN id, CAST(tags, 'STRING[3]') + """ + result = list(self.conn.execute(query)) + assert len(result) == 2 + assert result[0][1] == ["hello", "world", "test"] + assert result[1][1] == ["foo", "bar", "baz"] + + def test_cast_nested_int32_array(self): + """LOAD FROM CSV with CAST(col, 'INT32[2][2]') - nested fixed-size array.""" + csv_path = self._write_csv( + "nested_int_array.csv", + "id|matrix\n" "1|[[1, 2], [3, 4]]\n" "2|[[5, 6], [7, 8]]\n", + ) + query = f""" + LOAD FROM "{csv_path}" (delim='|') + RETURN id, CAST(matrix, 'INT32[2][2]') + """ + result = list(self.conn.execute(query)) + assert len(result) == 2 + assert result[0][1] == [[1, 2], [3, 4]] + assert result[1][1] == [[5, 6], [7, 8]] + + def test_cast_nested_string_array(self): + """LOAD FROM CSV with CAST(col, 'STRING[2][2]') - nested string array.""" + csv_path = self._write_csv( + "nested_string_array.csv", + "id|data\n" "1|[[a, b], [c, d]]\n" "2|[[w, x], [y, z]]\n", + ) + query = f""" + LOAD FROM "{csv_path}" (delim='|') + RETURN id, CAST(data, 'STRING[2][2]') + """ + result = list(self.conn.execute(query)) + assert len(result) == 2 + assert result[0][1] == [["a", "b"], ["c", "d"]] + assert result[1][1] == [["w", "x"], ["y", "z"]] + + def test_cast_date_array(self): + """LOAD FROM CSV with CAST(col, 'DATE[2]').""" + csv_path = self._write_csv( + "date_array.csv", + "id|dates\n" "1|[1970-01-01, 2023-06-15]\n" "2|[2000-01-01, 2001-01-01]\n", + ) + query = f""" + LOAD FROM "{csv_path}" (delim='|') + RETURN id, CAST(dates, 'DATE[2]') + """ + result = list(self.conn.execute(query)) + assert len(result) == 2 + assert len(result[0][1]) == 2 + # Verify dates are returned as datetime.date objects + assert result[0][1][0] == date(1970, 1, 1) + assert result[0][1][1] == date(2023, 6, 15) + assert len(result[1][1]) == 2 + assert result[1][1][0] == date(2000, 1, 1) + assert result[1][1][1] == date(2001, 1, 1) + + def test_cast_timestamp_array(self): + """LOAD FROM CSV with CAST(col, 'TIMESTAMP[2]').""" + csv_path = self._write_csv( + "timestamp_array.csv", + "id|times\n" + "1|[1970-01-01 00:00:00, 2023-06-15 12:30:00]\n" + "2|[2000-01-01 00:00:00, 2001-01-01 00:00:00]\n", + ) + query = f""" + LOAD FROM "{csv_path}" (delim='|') + RETURN id, CAST(times, 'TIMESTAMP[2]') + """ + result = list(self.conn.execute(query)) + assert len(result) == 2 + assert len(result[0][1]) == 2 + # Verify timestamps are returned as datetime.datetime objects + assert result[0][1][0] == datetime(1970, 1, 1, 0, 0) + assert result[0][1][1] == datetime(2023, 6, 15, 12, 30) + assert len(result[1][1]) == 2 + assert result[1][1][0] == datetime(2000, 1, 1, 0, 0) + assert result[1][1][1] == datetime(2001, 1, 1, 0, 0) + + def test_cast_interval_array(self): + """LOAD FROM CSV with CAST(col, 'INTERVAL[2]').""" + csv_path = self._write_csv( + "interval_array.csv", + "id|durations\n" "1|[2 days, 3 hours]\n" "2|[1 year 2 months, 4 days]\n", + ) + query = f""" + LOAD FROM "{csv_path}" (delim='|') + RETURN id, CAST(durations, 'INTERVAL[2]') + """ + result = list(self.conn.execute(query)) + assert len(result) == 2 + # Verify intervals are returned as strings + assert result[0][1] == ["2 days", "3 hours"] + assert result[1][1] == ["1 year 2 months", "4 days"] + + def test_cast_multiple_array_columns(self): + """LOAD FROM CSV with multiple CAST array columns.""" + csv_path = self._write_csv( + "multi_array.csv", + "id|ints|floats\n" "1|[1, 2, 3]|[1.1, 2.2]\n" "2|[4, 5, 6]|[3.3, 4.4]\n", + ) + query = f""" + LOAD FROM "{csv_path}" (delim='|') + RETURN id, CAST(ints, 'INT64[3]'), CAST(floats, 'DOUBLE[2]') + """ + result = list(self.conn.execute(query)) + assert len(result) == 2 + assert result[0][1] == [1, 2, 3] + assert len(result[0][2]) == 2 + assert abs(result[0][2][0] - 1.1) < 1e-9 + assert result[1][1] == [4, 5, 6] + + def test_cast_json_numeric_arrays(self): + """LOAD FROM JSON with numeric fixed-size array casts.""" + json_path = self._write_json( + "numeric_arrays.json", + [ + { + "id": 1, + "floats": [1.5, 2.5, 3.5], + "doubles": [1.111111111, 2.222222222], + "int32s": [10, 20], + "int64s": [100000000000, 200000000000], + }, + { + "id": 2, + "floats": [4.0, 5.0, 6.0], + "doubles": [3.333333333, 4.444444444], + "int32s": [-1, 0], + "int64s": [0, 1], + }, + ], + ) + query = f""" + LOAD FROM "{json_path}" + RETURN id, + CAST(floats, 'FLOAT[3]'), + CAST(doubles, 'DOUBLE[2]'), + CAST(int32s, 'INT32[2]'), + CAST(int64s, 'INT64[2]') + """ + result = list(self.conn.execute(query)) + assert len(result) == 2 + + assert len(result[0][1]) == 3 + assert abs(result[0][1][0] - 1.5) < 1e-5 + assert abs(result[0][1][1] - 2.5) < 1e-5 + assert abs(result[0][1][2] - 3.5) < 1e-5 + assert result[0][2] == [1.111111111, 2.222222222] + assert result[0][3] == [10, 20] + assert result[0][4] == [100000000000, 200000000000] + + assert len(result[1][1]) == 3 + assert abs(result[1][1][0] - 4.0) < 1e-5 + assert result[1][2] == [3.333333333, 4.444444444] + assert result[1][3] == [-1, 0] + assert result[1][4] == [0, 1] + + def test_cast_json_string_and_nested_arrays(self): + """LOAD FROM JSON with string and nested fixed-size array casts.""" + json_path = self._write_json( + "nested_arrays.json", + [ + { + "id": 1, + "tags": ["hello", "world"], + "matrix": [[1, 2], [3, 4]], + "labels": [["a", "b"], ["c", "d"]], + }, + { + "id": 2, + "tags": ["foo", "bar"], + "matrix": [[5, 6], [7, 8]], + "labels": [["w", "x"], ["y", "z"]], + }, + ], + ) + query = f""" + LOAD FROM "{json_path}" + RETURN id, + CAST(tags, 'STRING[2]'), + CAST(matrix, 'INT32[2][2]'), + CAST(labels, 'STRING[2][2]') + """ + result = list(self.conn.execute(query)) + assert len(result) == 2 + assert result[0][1] == ["hello", "world"] + assert result[0][2] == [[1, 2], [3, 4]] + assert result[0][3] == [["a", "b"], ["c", "d"]] + assert result[1][1] == ["foo", "bar"] + assert result[1][2] == [[5, 6], [7, 8]] + assert result[1][3] == [["w", "x"], ["y", "z"]] + + def test_cast_json_temporal_arrays(self): + """LOAD FROM JSON with temporal fixed-size array casts.""" + json_path = self._write_json( + "temporal_arrays.json", + [ + { + "id": 1, + "dates": ["1970-01-01", "2023-06-15"], + "times": ["1970-01-01 00:00:00", "2023-06-15 12:30:00"], + "durations": ["2 days", "3 hours"], + }, + { + "id": 2, + "dates": ["2000-01-01", "2001-01-01"], + "times": ["2000-01-01 00:00:00", "2001-01-01 00:00:00"], + "durations": ["1 year 2 months", "4 days"], + }, + ], + ) + query = f""" + LOAD FROM "{json_path}" + RETURN id, + CAST(dates, 'DATE[2]'), + CAST(times, 'TIMESTAMP[2]'), + CAST(durations, 'INTERVAL[2]') + """ + result = list(self.conn.execute(query)) + assert len(result) == 2 + + assert result[0][1] == [date(1970, 1, 1), date(2023, 6, 15)] + assert result[0][2] == [ + datetime(1970, 1, 1, 0, 0), + datetime(2023, 6, 15, 12, 30), + ] + assert result[0][3] == ["2 days", "3 hours"] + + assert result[1][1] == [date(2000, 1, 1), date(2001, 1, 1)] + assert result[1][2] == [ + datetime(2000, 1, 1, 0, 0), + datetime(2001, 1, 1, 0, 0), + ] + assert result[1][3] == ["1 year 2 months", "4 days"]