-
Notifications
You must be signed in to change notification settings - Fork 176
Expand file tree
/
Copy pathtest_context.py
More file actions
1895 lines (1437 loc) · 61.9 KB
/
Copy pathtest_context.py
File metadata and controls
1895 lines (1437 loc) · 61.9 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
580
581
582
583
584
585
586
587
588
589
590
591
592
593
594
595
596
597
598
599
600
601
602
603
604
605
606
607
608
609
610
611
612
613
614
615
616
617
618
619
620
621
622
623
624
625
626
627
628
629
630
631
632
633
634
635
636
637
638
639
640
641
642
643
644
645
646
647
648
649
650
651
652
653
654
655
656
657
658
659
660
661
662
663
664
665
666
667
668
669
670
671
672
673
674
675
676
677
678
679
680
681
682
683
684
685
686
687
688
689
690
691
692
693
694
695
696
697
698
699
700
701
702
703
704
705
706
707
708
709
710
711
712
713
714
715
716
717
718
719
720
721
722
723
724
725
726
727
728
729
730
731
732
733
734
735
736
737
738
739
740
741
742
743
744
745
746
747
748
749
750
751
752
753
754
755
756
757
758
759
760
761
762
763
764
765
766
767
768
769
770
771
772
773
774
775
776
777
778
779
780
781
782
783
784
785
786
787
788
789
790
791
792
793
794
795
796
797
798
799
800
801
802
803
804
805
806
807
808
809
810
811
812
813
814
815
816
817
818
819
820
821
822
823
824
825
826
827
828
829
830
831
832
833
834
835
836
837
838
839
840
841
842
843
844
845
846
847
848
849
850
851
852
853
854
855
856
857
858
859
860
861
862
863
864
865
866
867
868
869
870
871
872
873
874
875
876
877
878
879
880
881
882
883
884
885
886
887
888
889
890
891
892
893
894
895
896
897
898
899
900
901
902
903
904
905
906
907
908
909
910
911
912
913
914
915
916
917
918
919
920
921
922
923
924
925
926
927
928
929
930
931
932
933
934
935
936
937
938
939
940
941
942
943
944
945
946
947
948
949
950
951
952
953
954
955
956
957
958
959
960
961
962
963
964
965
966
967
968
969
970
971
972
973
974
975
976
977
978
979
980
981
982
983
984
985
986
987
988
989
990
991
992
993
994
995
996
997
998
999
1000
# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements. See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership. The ASF licenses this file
# to you 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.
import ctypes
import datetime as dt
import gc
import gzip
import pathlib
import shutil
import pyarrow as pa
import pyarrow.dataset as ds
import pytest
from datafusion import (
CsvReadOptions,
DataFrame,
RuntimeEnvBuilder,
SessionConfig,
SessionContext,
SessionExtensionComponents,
SQLOptions,
Table,
column,
literal,
udf,
)
def test_create_context_no_args():
SessionContext()
def test_create_context_session_config_only():
SessionContext(config=SessionConfig())
def test_create_context_runtime_config_only():
SessionContext(runtime=RuntimeEnvBuilder())
@pytest.mark.parametrize("path_to_str", [True, False])
def test_runtime_configs(tmp_path, path_to_str):
path1 = tmp_path / "dir1"
path2 = tmp_path / "dir2"
path1 = str(path1) if path_to_str else path1
path2 = str(path2) if path_to_str else path2
runtime = RuntimeEnvBuilder().with_disk_manager_specified(path1, path2)
config = SessionConfig().with_default_catalog_and_schema("foo", "bar")
ctx = SessionContext(config, runtime)
assert ctx is not None
db = ctx.catalog("foo").schema("bar")
assert db is not None
@pytest.mark.parametrize("path_to_str", [True, False])
def test_temporary_files(tmp_path, path_to_str):
path = str(tmp_path) if path_to_str else tmp_path
runtime = RuntimeEnvBuilder().with_temp_file_path(path)
config = SessionConfig().with_default_catalog_and_schema("foo", "bar")
ctx = SessionContext(config, runtime)
assert ctx is not None
db = ctx.catalog("foo").schema("bar")
assert db is not None
def test_create_context_with_all_valid_args():
runtime = RuntimeEnvBuilder().with_disk_manager_os().with_fair_spill_pool(10000000)
config = (
SessionConfig()
.with_create_default_catalog_and_schema(enabled=True)
.with_default_catalog_and_schema("foo", "bar")
.with_target_partitions(1)
.with_information_schema(enabled=True)
.with_repartition_joins(enabled=False)
.with_repartition_aggregations(enabled=False)
.with_repartition_windows(enabled=False)
.with_parquet_pruning(enabled=False)
)
ctx = SessionContext(config, runtime)
# verify that at least some of the arguments worked
ctx.catalog("foo").schema("bar")
with pytest.raises(KeyError):
ctx.catalog("datafusion")
def test_session_config_set_rejects_an_unknown_namespace():
"""A bad config key raises rather than aborting through a Rust panic.
`datafusion.runtime.*` appears in `information_schema.df_settings` but has
no `ConfigOptions` namespace, so it is the key a naive "read the settings
back and replay them on the worker" loop hits first.
"""
# `ValueError`, not a bare `Exception`: a panic would arrive as
# `PanicException`, which derives from `BaseException` and so would not be
# caught here at all. Both this and the constructor cases below rely on it.
with pytest.raises(ValueError, match="runtime"):
SessionConfig().set("datafusion.runtime.memory_limit", "unlimited")
def test_session_config_set_rejects_an_unparsable_value():
"""A well-known key with a value of the wrong type raises too."""
with pytest.raises(ValueError, match="batch_size"):
SessionConfig().set("datafusion.execution.batch_size", "not_an_int")
def test_session_config_constructor_applies_options():
"""A dict passed to the constructor reaches the session's options."""
config = SessionConfig(
{
"datafusion.execution.batch_size": "1024",
"datafusion.execution.target_partitions": "3",
}
)
ctx = SessionContext(config.with_information_schema(True))
settings = ctx.sql(
"select name, value from information_schema.df_settings"
" where name in ('datafusion.execution.batch_size',"
" 'datafusion.execution.target_partitions')"
).to_pydict()
assert dict(zip(settings["name"], settings["value"], strict=True)) == {
"datafusion.execution.batch_size": "1024",
"datafusion.execution.target_partitions": "3",
}
def test_session_config_constructor_rejects_an_unknown_namespace():
"""A bad key in the constructor's dict raises rather than panicking.
The same defect as `SessionConfig.set` had, reached through the argument
that a replayed `information_schema.df_settings` dictionary arrives in.
"""
with pytest.raises(ValueError, match="runtime"):
SessionConfig({"datafusion.runtime.memory_limit": "unlimited"})
def test_session_config_constructor_rejects_an_unparsable_value():
"""A well-known key with a value of the wrong type raises too."""
with pytest.raises(ValueError, match="batch_size"):
SessionConfig({"datafusion.execution.batch_size": "not_an_int"})
def test_register_record_batches(ctx):
# create a RecordBatch and register it as memtable
batch = pa.RecordBatch.from_arrays(
[pa.array([1, 2, 3]), pa.array([4, 5, 6])],
names=["a", "b"],
)
ctx.register_record_batches("t", [[batch]])
assert ctx.catalog().schema().names() == {"t"}
result = ctx.sql("SELECT a+b, a-b FROM t").collect()
assert result[0].column(0) == pa.array([5, 7, 9])
assert result[0].column(1) == pa.array([-3, -3, -3])
def test_register_record_batches_empty(ctx):
# A partition list with no record batches carries no schema, so this used to
# panic on unchecked `[0][0]` indexing. It should now raise a clear error.
with pytest.raises(ValueError, match="no record batches"):
ctx.register_record_batches("t", [[]])
# An empty outer partition list carries no schema either, and raises the same error.
with pytest.raises(ValueError, match="no record batches"):
ctx.register_record_batches("t", [])
# The schema is still recovered from a later non-empty partition.
batch = pa.RecordBatch.from_arrays([pa.array([1, 2, 3])], names=["a"])
ctx.register_record_batches("t2", [[], [batch]])
assert ctx.sql("SELECT a FROM t2").collect()[0].column(0) == pa.array([1, 2, 3])
def test_create_dataframe_registers_unique_table_name(ctx):
# create a RecordBatch and register it as memtable
batch = pa.RecordBatch.from_arrays(
[pa.array([1, 2, 3]), pa.array([4, 5, 6])],
names=["a", "b"],
)
df = ctx.create_dataframe([[batch]])
tables = list(ctx.catalog().schema().names())
assert df
assert len(tables) == 1
assert len(tables[0]) == 33
assert tables[0].startswith("c")
# ensure that the rest of the table name contains
# only hexadecimal numbers
for c in tables[0][1:]:
assert c in "0123456789abcdef"
def test_create_dataframe_registers_with_defined_table_name(ctx):
# create a RecordBatch and register it as memtable
batch = pa.RecordBatch.from_arrays(
[pa.array([1, 2, 3]), pa.array([4, 5, 6])],
names=["a", "b"],
)
df = ctx.create_dataframe([[batch]], name="tbl")
tables = list(ctx.catalog().schema().names())
assert df
assert len(tables) == 1
assert tables[0] == "tbl"
def test_from_arrow_table(ctx):
# create a PyArrow table
data = {"a": [1, 2, 3], "b": [4, 5, 6]}
table = pa.Table.from_pydict(data)
# convert to DataFrame
df = ctx.from_arrow(table)
tables = list(ctx.catalog().schema().names())
assert df
assert len(tables) == 1
assert isinstance(df, DataFrame)
assert set(df.schema().names) == {"a", "b"}
assert df.collect()[0].num_rows == 3
def record_batch_generator(num_batches: int):
schema = pa.schema([("a", pa.int64()), ("b", pa.int64())])
for _i in range(num_batches):
yield pa.RecordBatch.from_arrays(
[pa.array([1, 2, 3]), pa.array([4, 5, 6])], schema=schema
)
@pytest.mark.parametrize(
"source",
[
# __arrow_c_array__ sources
pa.array([{"a": 1, "b": 4}, {"a": 2, "b": 5}, {"a": 3, "b": 6}]),
# __arrow_c_stream__ sources
pa.RecordBatch.from_pydict({"a": [1, 2, 3], "b": [4, 5, 6]}),
pa.RecordBatchReader.from_batches(
pa.schema([("a", pa.int64()), ("b", pa.int64())]), record_batch_generator(1)
),
pa.Table.from_pydict({"a": [1, 2, 3], "b": [4, 5, 6]}),
],
)
def test_from_arrow_sources(ctx, source) -> None:
df = ctx.from_arrow(source)
assert df
assert isinstance(df, DataFrame)
assert df.schema().names == ["a", "b"]
assert df.count() == 3
def test_from_arrow_table_with_name(ctx):
# create a PyArrow table
data = {"a": [1, 2, 3], "b": [4, 5, 6]}
table = pa.Table.from_pydict(data)
# convert to DataFrame with optional name
df = ctx.from_arrow(table, name="tbl")
tables = list(ctx.catalog().schema().names())
assert df
assert tables[0] == "tbl"
def test_from_arrow_table_empty(ctx):
data = {"a": [], "b": []}
schema = pa.schema([("a", pa.int32()), ("b", pa.string())])
table = pa.Table.from_pydict(data, schema=schema)
# convert to DataFrame
df = ctx.from_arrow(table)
tables = list(ctx.catalog().schema().names())
assert df
assert len(tables) == 1
assert isinstance(df, DataFrame)
assert set(df.schema().names) == {"a", "b"}
assert len(df.collect()) == 0
def test_from_arrow_table_empty_no_schema(ctx):
data = {"a": [], "b": []}
table = pa.Table.from_pydict(data)
# convert to DataFrame
df = ctx.from_arrow(table)
tables = list(ctx.catalog().schema().names())
assert df
assert len(tables) == 1
assert isinstance(df, DataFrame)
assert set(df.schema().names) == {"a", "b"}
assert len(df.collect()) == 0
def test_from_pylist(ctx):
# create a dataframe from Python list
data = [
{"a": 1, "b": 4},
{"a": 2, "b": 5},
{"a": 3, "b": 6},
]
df = ctx.from_pylist(data)
tables = list(ctx.catalog().schema().names())
assert df
assert len(tables) == 1
assert isinstance(df, DataFrame)
assert set(df.schema().names) == {"a", "b"}
assert df.collect()[0].num_rows == 3
def test_from_pydict(ctx):
# create a dataframe from Python dictionary
data = {"a": [1, 2, 3], "b": [4, 5, 6]}
df = ctx.from_pydict(data)
tables = list(ctx.catalog().schema().names())
assert df
assert len(tables) == 1
assert isinstance(df, DataFrame)
assert set(df.schema().names) == {"a", "b"}
assert df.collect()[0].num_rows == 3
def test_from_pandas(ctx):
# create a dataframe from pandas dataframe
pd = pytest.importorskip("pandas")
data = {"a": [1, 2, 3], "b": [4, 5, 6]}
pandas_df = pd.DataFrame(data)
df = ctx.from_pandas(pandas_df)
tables = list(ctx.catalog().schema().names())
assert df
assert len(tables) == 1
assert isinstance(df, DataFrame)
assert set(df.schema().names) == {"a", "b"}
assert df.collect()[0].num_rows == 3
def test_from_polars(ctx):
# create a dataframe from Polars dataframe
pd = pytest.importorskip("polars")
data = {"a": [1, 2, 3], "b": [4, 5, 6]}
polars_df = pd.DataFrame(data)
df = ctx.from_polars(polars_df)
tables = list(ctx.catalog().schema().names())
assert df
assert len(tables) == 1
assert isinstance(df, DataFrame)
assert set(df.schema().names) == {"a", "b"}
assert df.collect()[0].num_rows == 3
def test_register_table(ctx, database):
default = ctx.catalog()
public = default.schema("public")
assert public.names() == {"csv", "csv1", "csv2"}
table = public.table("csv")
ctx.register_table("csv3", table)
assert public.names() == {"csv", "csv1", "csv2", "csv3"}
def test_read_table_from_catalog(ctx, database):
default = ctx.catalog()
public = default.schema("public")
assert public.names() == {"csv", "csv1", "csv2"}
table = public.table("csv")
table_df = ctx.read_table(table)
table_df.show()
def test_read_table_from_df(ctx):
df = ctx.from_pydict({"a": [1, 2]})
result = ctx.read_table(df).collect()
assert [b.to_pydict() for b in result] == [{"a": [1, 2]}]
def test_read_table_from_dataset(ctx):
batch = pa.RecordBatch.from_arrays(
[pa.array([1, 2, 3]), pa.array([4, 5, 6])],
names=["a", "b"],
)
dataset = ds.dataset([batch])
result = ctx.read_table(dataset).collect()
assert result[0].column(0) == pa.array([1, 2, 3])
assert result[0].column(1) == pa.array([4, 5, 6])
def test_deregister_table(ctx, database):
default = ctx.catalog()
public = default.schema("public")
assert public.names() == {"csv", "csv1", "csv2"}
ctx.deregister_table("csv")
assert public.names() == {"csv1", "csv2"}
def test_deregister_udf():
ctx = SessionContext()
is_null = udf(
lambda x: x.is_null(),
[pa.float64()],
pa.bool_(),
volatility="immutable",
name="my_is_null",
)
ctx.register_udf(is_null)
# Verify it works
df = ctx.from_pydict({"a": [1.0, None]})
ctx.register_table("t", df.into_view())
result = ctx.sql("SELECT my_is_null(a) FROM t").collect()
assert result[0].column(0) == pa.array([False, True])
# Deregister and verify it's gone
ctx.deregister_udf("my_is_null")
with pytest.raises(ValueError):
ctx.sql("SELECT my_is_null(a) FROM t").collect()
def test_deregister_udaf():
import pyarrow.compute as pc
ctx = SessionContext()
from datafusion import Accumulator, udaf
class MySum(Accumulator):
def __init__(self):
self._sum = 0.0
def update(self, values: pa.Array) -> None:
self._sum += pc.sum(values).as_py()
def merge(self, states: list[pa.Array]) -> None:
self._sum += pc.sum(states[0]).as_py()
def state(self) -> list:
return [self._sum]
def evaluate(self) -> pa.Scalar:
return self._sum
my_sum = udaf(
MySum,
[pa.float64()],
pa.float64(),
[pa.float64()],
volatility="immutable",
name="my_sum",
)
ctx.register_udaf(my_sum)
df = ctx.from_pydict({"a": [1.0, 2.0, 3.0]})
ctx.register_table("t", df.into_view())
result = ctx.sql("SELECT my_sum(a) FROM t").collect()
assert result[0].column(0) == pa.array([6.0])
ctx.deregister_udaf("my_sum")
with pytest.raises(ValueError):
ctx.sql("SELECT my_sum(a) FROM t").collect()
def test_deregister_udwf():
ctx = SessionContext()
from datafusion import udwf
from datafusion.user_defined import WindowEvaluator
class MyRowNumber(WindowEvaluator):
def __init__(self):
self._row = 0
def evaluate_all(self, values, num_rows):
return pa.array(list(range(1, num_rows + 1)), type=pa.uint64())
my_row_number = udwf(
MyRowNumber,
[pa.float64()],
pa.uint64(),
volatility="immutable",
name="my_row_number",
)
ctx.register_udwf(my_row_number)
df = ctx.from_pydict({"a": [1.0, 2.0, 3.0]})
ctx.register_table("t", df.into_view())
result = ctx.sql("SELECT my_row_number(a) OVER () FROM t").collect()
assert result[0].column(0) == pa.array([1, 2, 3], type=pa.uint64())
ctx.deregister_udwf("my_row_number")
with pytest.raises(ValueError):
ctx.sql("SELECT my_row_number(a) OVER () FROM t").collect()
def test_deregister_udtf():
import pyarrow.dataset as ds
ctx = SessionContext()
from datafusion import Table, udtf
class MyTable:
def __call__(self):
batch = pa.RecordBatch.from_pydict({"x": [1, 2, 3]})
return Table(ds.dataset([batch]))
my_table = udtf(MyTable(), "my_table")
ctx.register_udtf(my_table)
result = ctx.sql("SELECT * FROM my_table()").collect()
assert result[0].column(0) == pa.array([1, 2, 3])
ctx.deregister_udtf("my_table")
with pytest.raises(ValueError):
ctx.sql("SELECT * FROM my_table()").collect()
def test_register_table_from_dataframe(ctx):
df = ctx.from_pydict({"a": [1, 2]})
ctx.register_table("df_tbl", df)
result = ctx.sql("SELECT * FROM df_tbl").collect()
assert [b.to_pydict() for b in result] == [{"a": [1, 2]}]
@pytest.mark.parametrize("temporary", [True, False])
def test_register_table_from_dataframe_into_view(ctx, temporary):
df = ctx.from_pydict({"a": [1, 2]})
table = df.into_view(temporary=temporary)
assert isinstance(table, Table)
if temporary:
assert table.kind == "temporary"
else:
assert table.kind == "view"
ctx.register_table("view_tbl", table)
result = ctx.sql("SELECT * FROM view_tbl").collect()
assert [b.to_pydict() for b in result] == [{"a": [1, 2]}]
def test_table_from_dataframe(ctx):
df = ctx.from_pydict({"a": [1, 2]})
table = Table(df)
assert isinstance(table, Table)
ctx.register_table("from_dataframe_tbl", table)
result = ctx.sql("SELECT * FROM from_dataframe_tbl").collect()
assert [b.to_pydict() for b in result] == [{"a": [1, 2]}]
def test_table_from_dataframe_internal(ctx):
df = ctx.from_pydict({"a": [1, 2]})
table = Table(df.df)
assert isinstance(table, Table)
ctx.register_table("from_internal_dataframe_tbl", table)
result = ctx.sql("SELECT * FROM from_internal_dataframe_tbl").collect()
assert [b.to_pydict() for b in result] == [{"a": [1, 2]}]
def test_register_dataset(ctx):
# create a RecordBatch and register it as a pyarrow.dataset.Dataset
batch = pa.RecordBatch.from_arrays(
[pa.array([1, 2, 3]), pa.array([4, 5, 6])],
names=["a", "b"],
)
dataset = ds.dataset([batch])
ctx.register_dataset("t", dataset)
assert ctx.catalog().schema().names() == {"t"}
result = ctx.sql("SELECT a+b, a-b FROM t").collect()
assert result[0].column(0) == pa.array([5, 7, 9])
assert result[0].column(1) == pa.array([-3, -3, -3])
def test_dataset_filter(ctx, capfd):
# create a RecordBatch and register it as a pyarrow.dataset.Dataset
batch = pa.RecordBatch.from_arrays(
[pa.array([1, 2, 3]), pa.array([4, 5, 6])],
names=["a", "b"],
)
dataset = ds.dataset([batch])
ctx.register_dataset("t", dataset)
assert ctx.catalog().schema().names() == {"t"}
df = ctx.sql("SELECT a+b, a-b FROM t WHERE a BETWEEN 2 and 3 AND b > 5")
# Make sure the filter was pushed down in Physical Plan
df.explain()
captured = capfd.readouterr()
assert "filter_expr=(((a >= 2) and (a <= 3)) and (b > 5))" in captured.out
result = df.collect()
assert result[0].column(0) == pa.array([9])
assert result[0].column(1) == pa.array([-3])
def test_dataset_count(ctx):
# `datafusion-python` issue: https://github.com/apache/datafusion-python/issues/800
batch = pa.RecordBatch.from_arrays(
[pa.array([1, 2, 3]), pa.array([4, 5, 6])],
names=["a", "b"],
)
dataset = ds.dataset([batch])
ctx.register_dataset("t", dataset)
# Testing the dataframe API
df = ctx.table("t")
assert df.count() == 3
# Testing the SQL API
count = ctx.sql("SELECT COUNT(*) FROM t")
count = count.collect()
assert count[0].column(0) == pa.array([3])
def test_pyarrow_predicate_pushdown_is_null(ctx, capfd):
"""Ensure that pyarrow filter gets pushed down for `IsNull`"""
# create a RecordBatch and register it as a pyarrow.dataset.Dataset
batch = pa.RecordBatch.from_arrays(
[pa.array([1, 2, 3]), pa.array([4, 5, 6]), pa.array([7, None, 9])],
names=["a", "b", "c"],
)
dataset = ds.dataset([batch])
ctx.register_dataset("t", dataset)
# Make sure the filter was pushed down in Physical Plan
df = ctx.sql("SELECT a FROM t WHERE c is NULL")
df.explain()
captured = capfd.readouterr()
assert "filter_expr=is_null(c, {nan_is_null=false})" in captured.out
result = df.collect()
assert result[0].column(0) == pa.array([2])
def test_pyarrow_predicate_pushdown_timestamp(ctx, tmpdir, capfd):
"""Ensure that pyarrow filter gets pushed down for timestamp"""
# Ref: https://github.com/apache/datafusion-python/issues/703
# create pyarrow dataset with no actual files
col_type = pa.timestamp("ns", "+00:00")
nyd_2000 = pa.scalar(dt.datetime(2000, 1, 1, tzinfo=dt.timezone.utc), col_type)
pa_dataset_fs = pa.fs.SubTreeFileSystem(str(tmpdir), pa.fs.LocalFileSystem())
pa_dataset_format = pa.dataset.ParquetFileFormat()
pa_dataset_partition = pa.dataset.field("a") <= nyd_2000
fragments = [
# NOTE: we never actually make this file.
# Working predicate pushdown means it never gets accessed
pa_dataset_format.make_fragment(
"1.parquet",
filesystem=pa_dataset_fs,
partition_expression=pa_dataset_partition,
)
]
pa_dataset = pa.dataset.FileSystemDataset(
fragments,
pa.schema([pa.field("a", col_type)]),
pa_dataset_format,
pa_dataset_fs,
)
ctx.register_dataset("t", pa_dataset)
# the partition for our only fragment is for a < 2000-01-01.
# so querying for a > 2024-01-01 should not touch any files
df = ctx.sql("SELECT * FROM t WHERE a > '2024-01-01T00:00:00+00:00'")
assert df.collect() == []
def test_dataset_filter_nested_data(ctx):
# create Arrow StructArrays to test nested data types
data = pa.StructArray.from_arrays(
[pa.array([1, 2, 3]), pa.array([4, 5, 6])],
names=["a", "b"],
)
batch = pa.RecordBatch.from_arrays(
[data],
names=["nested_data"],
)
dataset = ds.dataset([batch])
ctx.register_dataset("t", dataset)
assert ctx.catalog().schema().names() == {"t"}
df = ctx.table("t")
# This filter will not be pushed down to DatasetExec since it
# isn't supported
df = df.filter(column("nested_data")["b"] > literal(5)).select(
column("nested_data")["a"] + column("nested_data")["b"],
column("nested_data")["a"] - column("nested_data")["b"],
)
result = df.collect()
assert result[0].column(0) == pa.array([9])
assert result[0].column(1) == pa.array([-3])
def test_table_exist(ctx):
batch = pa.RecordBatch.from_arrays(
[pa.array([1, 2, 3]), pa.array([4, 5, 6])],
names=["a", "b"],
)
dataset = ds.dataset([batch])
ctx.register_dataset("t", dataset)
assert ctx.table_exist("t") is True
def test_table_not_found(ctx):
from uuid import uuid4
with pytest.raises(KeyError):
ctx.table(f"not-found-{uuid4()}")
def test_session_start_time(ctx):
import datetime
import re
st = ctx.session_start_time()
assert isinstance(st, str)
# Truncate nanoseconds to microseconds for Python 3.10 compat
st = re.sub(r"(\.\d{6})\d+", r"\1", st)
dt = datetime.datetime.fromisoformat(st)
assert dt.isoformat()
def test_enable_ident_normalization(ctx):
assert ctx.enable_ident_normalization() is True
ctx.sql("SET datafusion.sql_parser.enable_ident_normalization = false")
assert ctx.enable_ident_normalization() is False
def test_parse_sql_expr(ctx):
from datafusion.common import DFSchema
schema = DFSchema.empty()
expr = ctx.parse_sql_expr("1 + 2", schema)
assert str(expr) == "Expr(Int64(1) + Int64(2))"
def test_execute_logical_plan(ctx):
df = ctx.from_pydict({"a": [1, 2, 3]})
plan = df.logical_plan()
df2 = ctx.execute_logical_plan(plan)
result = df2.collect()
assert result[0].column(0) == pa.array([1, 2, 3])
def test_refresh_catalogs(ctx):
ctx.refresh_catalogs()
def test_remove_optimizer_rule(ctx):
assert ctx.remove_optimizer_rule("push_down_filter") is True
assert ctx.remove_optimizer_rule("nonexistent_rule") is False
def test_set_query_planner_rejects_wrong_capsule(ctx):
with pytest.raises(ValueError, match="datafusion_query_planner"):
ctx.set_query_planner(ctx.__datafusion_task_context_provider__())
def test_with_extension_rejects_wrong_capsule(ctx):
"""The extension options hook names the capsule it was handed.
Like the rest of the capsule family, this reports which capsule turned up
rather than CPython's fixed "called with incorrect name" string.
"""
class WrongCapsule:
def __datafusion_extension_options__(self):
return ctx.__datafusion_task_context_provider__()
with pytest.raises(ValueError, match="datafusion_extension_options"):
SessionConfig().with_extension(WrongCapsule())
def test_pre_55_codec_signature_reports_an_upgrade(ctx):
"""A getter that refuses the session is named, not left as a bare TypeError.
Extension libraries implement these getters, so the pre-55.0.0 signature
is what an out-of-date one still has. The original error stays reachable
as ``__cause__`` rather than being replaced outright.
"""
class PreSessionCodec:
def __datafusion_logical_extension_codec__(self):
msg = "should never be called"
raise AssertionError(msg)
with pytest.raises(ImportError, match="__datafusion_logical_extension_codec__"):
ctx.with_logical_extension_codec(PreSessionCodec())
with pytest.raises(ImportError) as excinfo:
ctx.with_logical_extension_codec(PreSessionCodec())
assert isinstance(excinfo.value.__cause__, TypeError)
assert "positional argument" in str(excinfo.value.__cause__)
def test_type_error_inside_a_getter_is_not_reported_as_an_upgrade(ctx):
"""A correctly-signed getter's own TypeError must survive unchanged.
Only the call machinery's arity error means the library is out of date.
Rewriting every TypeError would send an author debugging their own getter
off to upgrade a library that is already correct.
"""
class RaisesTypeError:
def __datafusion_logical_extension_codec__(self, session):
msg = "bad cast inside the getter"
raise TypeError(msg)
with pytest.raises(TypeError, match="bad cast inside the getter"):
ctx.with_logical_extension_codec(RaisesTypeError())
def test_non_type_errors_from_a_getter_propagate(ctx):
"""Anything that is not a TypeError was never a signature problem."""
class RaisesValueError:
def __datafusion_logical_extension_codec__(self, session):
msg = "something else entirely"
raise ValueError(msg)
with pytest.raises(ValueError, match="something else entirely"):
ctx.with_logical_extension_codec(RaisesValueError())
def test_set_query_planner_capsule(ctx):
capsule = ctx.__datafusion_query_planner__()
get_name = ctypes.pythonapi.PyCapsule_GetName
get_name.argtypes = [ctypes.py_object]
get_name.restype = ctypes.c_char_p
assert get_name(capsule) == b"datafusion_query_planner"
ctx.register_record_batches(
"query_planner_test",
[[pa.RecordBatch.from_pydict({"value": [1, 2, 3]})]],
)
ctx.set_query_planner(capsule)
assert ctx.table_exist("query_planner_test")
batches = ctx.sql("SELECT 1 AS value").collect()
assert batches[0].column(0) == pa.array([1])
def test_installing_a_planner_leaves_the_session_intact(ctx):
"""The planner is written into the existing session, not a copy of it.
Registrations made before the install are still visible afterwards, and
ones made after are visible too -- there is a single session throughout,
so neither the catalogs nor the function registry are snapshotted.
"""
ctx.register_record_batches(
"registered_before",
[[pa.RecordBatch.from_pydict({"value": [1]})]],
)
before = udf(
lambda arr: arr,
[pa.int64()],
pa.int64(),
volatility="immutable",
name="registered_before",
)
ctx.register_udf(before)
ctx.set_query_planner(ctx.__datafusion_query_planner__())
ctx.register_record_batches(
"registered_after",
[[pa.RecordBatch.from_pydict({"value": [2]})]],
)
after = udf(
lambda arr: arr,
[pa.int64()],
pa.int64(),
volatility="immutable",
name="registered_after",
)
ctx.register_udf(after)
assert ctx.table_exist("registered_before")
assert ctx.table_exist("registered_after")
assert ctx.sql("SELECT registered_before(1)").collect()
assert ctx.sql("SELECT registered_after(1)").collect()
def test_contexts_sharing_a_session_share_the_planner(ctx):
"""A context derived before the install still plans through the planner.
``with_python_udf_inlining`` returns a handle on the same session, and the
query planner lives in that session's state.
"""
sibling = ctx.with_python_udf_inlining(enabled=False)
ctx.register_record_batches(
"shared_planner_test",
[[pa.RecordBatch.from_pydict({"value": [1, 2, 3]})]],
)
ctx.set_query_planner(ctx.__datafusion_query_planner__())
assert sibling.table_exist("shared_planner_test")
assert sibling.session_id() == ctx.session_id()
class _NamedCodec:
"""Wraps a codec capsule in an object that can name itself.
``with_extensions`` requires objects rather than bare capsules, because a
codec's wire id is read off the object it is handed over as. This is the
shape a library holding a raw capsule hands over.
"""
def __init__(self, capsule, codec_id):
self._capsule = capsule
self.__datafusion_codec_id__ = codec_id
def __datafusion_logical_extension_codec__(self, session=None):
return self._capsule
def __datafusion_physical_extension_codec__(self, session=None):
return self._capsule
class _CodecOnlyExtension:
"""Contributes decline-all codecs exported from an unrelated session.
Retaining ``ctx`` is what the protocol tells real extensions not to do —
a bundle is reusable, so a cached context belongs to whichever session it
was last installed on. It is kept here only so a test can assert *which*
context the factory was handed.
"""
def __init__(self, prefix="my_library"):
self.exporter = SessionContext()
self.prefix = prefix
self.bound_ctx = None
def __datafusion_session_components__(self, ctx):
self.bound_ctx = ctx
return SessionExtensionComponents(
logical_extension_codecs=(
_NamedCodec(
self.exporter.__datafusion_logical_extension_codec__(),
f"{self.prefix}.logical",
),
),
physical_extension_codecs=(
_NamedCodec(
self.exporter.__datafusion_physical_extension_codec__(),
f"{self.prefix}.physical",
),
),
)
class _PlannerExtension:
"""Contributes a planner, recording the fallback it was handed.
Passing ``fallback`` straight back through is the degenerate wrap: it plans
the same queries to the same plans, which is what lets a pure-Python test
assert the threading without a real layering planner. It is not a no-op —
the capsule gets installed, so the session ends up planning through a
foreign planner — but nothing here depends on that either way.